#%% Connect and set up SAS
import pandas as pd
import seaborn as sns
import logging
import time
from matplotlib import pyplot as plt

# Configure logging
log_file = rf"../logs/3_4_restatement_car_plots.txt"
logging.basicConfig(
    filename=log_file, 
    level=logging.INFO, 
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

def log(msg):
    print(msg)
    logging.info(msg)

def print_time(start):
    secs = time.time() - start
    s = int(secs % 60)
    m = int(secs / 60 % 60)
    h = int(secs / 3600)
    return f"{h:0>2}:{m:0>2}:{s:0>2}"

start = time.time()

# Read Data
try:
    df = pd.read_sas(r'../data/gl_returns.sas7bdat')
    log("SAS data loaded successfully.")
except Exception as e:
    log(f"ERROR reading SAS file: {e}")
    raise

# Convert specific columns to int, if they exist
cols_to_convert = ['bearish_engage_q5','bearish_engage_2_q5', 
                   'bearish_priority_q5', 'prin_sir_q5']
for c in cols_to_convert:
    if c in df.columns:
        df[c] = df[c].astype(int)
    else:
        log(f"WARNING: Column {c} not found in dataframe.")

#%% Plot Figure 1
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300
sns.set(font='Times New Roman', style='white')

fig, axs = plt.subplots(4, 1, figsize=(7, 10), sharey=True, constrained_layout=True)

sns.lineplot(x='t', y='car', data=df.loc[df.bearish_engage_q5.isin([1, 3, 5])], style='bearish_engage_q5', ax=axs[0], errorbar=None)
sns.lineplot(x='t', y='car', data=df.loc[df.bearish_engage_2_q5.isin([1, 3, 5])], style='bearish_engage_2_q5', ax=axs[1], errorbar=None)
sns.lineplot(x='t', y='car', data=df.loc[df.bearish_priority_q5.isin([1, 3, 5])], style='bearish_priority_q5', ax=axs[2], errorbar=None)
sns.lineplot(x='t', y='car', data=df.loc[df.prin_sir_q5.isin([1, 3, 5])], style='prin_sir_q5', ax=axs[3], errorbar=None)

for ax in axs.flatten():
    ax.set_title('', color='black')
    ax.set_xlabel('', color='black')
    ax.set_ylabel('', color='black')
    ax.axvline(x=0, color='red', linewidth=.5, ls=':')
    ax.legend(title='', loc='lower left')

fig.suptitle("Figure 1: CAR Around Peers' Restatement Announcements", fontweight='bold', fontsize=12)
axs[0].set_title("Panel A: by Quintiles of Engagement", loc='left', fontsize=12, fontweight='bold')
axs[1].set_title("Panel B: by Quintiles of Engagement %", loc='left', fontsize=12, fontweight='bold')
axs[2].set_title("Panel C: by Quintiles of Priority", loc='left', fontsize=12, fontweight='bold')

plt.savefig("../results/car.png")
log("Saved figure using the individual measures of Engagement to ../results/car.png")

# %% Plot Figure 2
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['font.family'] = "Times New Roman"
plt.rcParams['font.size'] = 12
plt.rcParams['text.usetex'] = False
sns.set(font='Times New Roman', style='white')

fig, ax = plt.subplots(1, 1, figsize=(7, 5), sharey=True, constrained_layout=True)

# Recode prin_sir_q5_2 for grouped lineplot
df['prin_sir_q5_2'] = df['prin_sir_q5']
df.loc[df['prin_sir_q5_2'].isin([2, 3, 4]), 'prin_sir_q5_2'] = 2
df.loc[df['prin_sir_q5_2'] == 5, 'prin_sir_q5_2'] = 3

sns.lineplot(x='t', y='car', data=df, style='prin_sir_q5_2', errorbar=None)

ax.axvline(x=0, color='red', linewidth=.5, ls=':')
h, l = ax.get_legend_handles_labels()

# Clean up legend
legend_map = {'1': 'Low', '2': 'Medium', '3': 'High'}
l = [legend_map.get(str(label), label) for label in l]
ax.legend(h, l, title='Engagement with Unfavorable Analysts', loc='upper left', ncol=3, labelcolor='black')

ax.set_xlabel("Trading Days Since Peer’s Restatement Announcement", color='black', fontsize=12)
ax.set_title("")
ax.set_ylabel("CAR", color='black', fontsize=12)

plt.savefig("../results/car2.png")
log("Saved figure using quintiles of PCF Enagement to ../results/car2.png")

# %%
log(f"Total runtime: {print_time(start)}")
log("All plots generated successfully.")