sqlPath = r'C:/Users/jaredflake/Dropbox/Research/Data/CIQ_Transcripts/2024/output/wrds_ciq_2024.db'
ciqPath = r"C:/Users/jaredflake/Dropbox/Research/Data/CIQ_Transcripts/2024"

import pandas as pd, os, time, sys, logging, numpy as np

sys.path.append(rf'modules')

import conversions as cv
from sqlite_addons import create_connection

# Configure logging
log_file = rf"../logs/0_0_classify_prep_log.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()
pd.set_option("display.precision", 3, 'display.float_format', '{:0.3f}'.format)


# Read Data
try:
    events = pd.read_pickle(os.path.join(ciqPath,"output/wrds_transcript_detail.pkl"), compression='zip')
    log("Successfully read wrds_transcript_detail.")
except Exception as e:
    log(f"Error reading wrds_transcript_detail: {e}")
    sys.exit(1)

try:
    wrds_gvkey = cv.readsas(os.path.join(ciqPath,"output/wrds_gvkey.sas7bdat"))
    log("Successfully read wrds_gvkey.")
except Exception as e:
    log(f"Error reading wrds_gvkey: {e}")
    sys.exit(1)

# Merge and filter data
log(f'events dataset has {len(events)} rows before filtering.')
try:
    events = events.merge(wrds_gvkey, 'left', on='companyid')
    del wrds_gvkey
    log("Merged wrds_gvkey into events dataset.")
except Exception as e:
    log(f"Error merging wrds_gvkey: {e}")
    sys.exit(1)

# Restricting event types
log("Filtering events dataset based on conditions.")
event_types = {'CC': 48, 'Guide_Update': 49, 'Investor_Analyst': 50, 'Conf': 51, 'M_A': 52, 'AI_Day': 192}
reqs = (events.companyid == events.companyid) & (events.transcriptpresentationtypename == "Final") & (events.keydeveventtypeid.isin([event_types['CC']]))
events = events[reqs].sort_values(['companyid', 'keydevid', 'transcriptcreationdate_utc', 'transcriptcreationtime_utc'], ascending=[True, True, False, False])
events = events.groupby(['companyid', 'keydevid']).nth(0).reset_index()
log(f'Filtered events dataset now has {len(events)} rows.')

# Processing company IDs
to_process = ['companyid', 'gvkey']
for col in to_process:
    start_time = time.time()
    if type(events[col]) != object:
        events[col] = events[col].astype(str)
    events[col] = events.groupby(['transcriptid'])[col].transform(' | '.join)
    comp = events[col].str.split(r'\s+\|\s+', expand=True).add_prefix(col)
    events = events.join(comp)
    del comp
    log(f"Processed {col} in {print_time(start_time)}")

events = events.drop_duplicates().drop(['companyid', 'gvkey'], axis=1).reset_index()
transcriptids = events.transcriptid.unique()
transcriptids = np.array([int(d) for d in transcriptids])

log("Dropped duplicates and reset index.")

# Read transcripts
try:
    conn = create_connection(sqlPath)
    transcripts = pd.read_sql_query(f"SELECT * FROM wrds_transcript_person WHERE transcriptid IN {tuple(transcriptids)}", conn)
    conn.close()
    log("Successfully read wrds_transcript_person.")
except Exception as e:
    log(f"Error reading wrds_transcript_person: {e}")
    sys.exit(1)

# Merge and save
try:
    transcripts = transcripts.merge(events, 'inner', on='transcriptid').drop_duplicates()
    transcripts.to_pickle(rf'../data/transcripts.pkl', compression='zip')
    log("Merged transcripts and events, and saved transcripts.pkl.")
except Exception as e:
    log(f"Error processing transcripts: {e}")
    sys.exit(1)

# Read and merge component text
timestamp = time.time()
transcriptids = transcripts.transcriptid.unique()
log(f'Reading component text for {len(transcriptids)} transcripts.')
try:
    conn = create_connection(sqlPath)
    df = pd.read_sql_query(f"SELECT transcriptid, transcriptcomponentid, componenttext FROM ciqtranscriptcomponent WHERE transcriptid IN {tuple(transcriptids)}", conn)
    conn.close()
    log(f'Read {len(df)} transcript components in {print_time(timestamp)}.')
except Exception as e:
    log(f"Error reading transcript components: {e}")
    sys.exit(1)

# Process Q&A
timestamp = time.time()
df = df.merge(transcripts[['transcriptid', 'transcriptcomponentid', 'transcriptcomponenttypeid', 'word_count']], how='left', on=['transcriptid', 'transcriptcomponentid'])
Q_A = df[df.transcriptcomponenttypeid.isin([3, 4, 8])].copy().sort_values(['transcriptid', 'transcriptcomponentid'])
Q_A['group_id'] = Q_A['transcriptcomponenttypeid'].eq(3).cumsum()
questions = Q_A[Q_A.transcriptcomponenttypeid == 3].copy().reset_index()
answers = Q_A[Q_A.transcriptcomponenttypeid.isin([4, 8])].dropna(subset=['componenttext'])
del Q_A, df
log(f'Processed Q&A in {print_time(timestamp)}.')

# Save files
timestamp = time.time()
try:
    questions.to_pickle(rf'../data/questions.pkl', compression='zip')
    answers.to_pickle(rf'../data/answers.pkl', compression='zip')
    log(f'Saved questions and answers in {print_time(timestamp)}.')
except Exception as e:
    log(f"Error saving questions/answers: {e}")
    sys.exit(1)

log(f"Script execution completed in {print_time(start)}.")
