import pandas as pd
import os

os.chdir(r'../Data/Data_website/Scrapping')

def uniqueURLgen(df):
    df = df[['urlkey', 'statuscode', 'length', 'datetime', 'year', 'baseKey', 'url', 'baseKey4Yr', 'fm_id']]
    df['date'] = df['datetime'].dt.date
    df['url_id'] = df.groupby(['urlkey']).ngroup()
    df[['statuscode','length']] = df[['statuscode','length']].apply(pd.to_numeric, errors='coerce', axis=1)
    df = df.sort_values('datetime').groupby(['urlkey','date']).tail(1).reset_index(drop=True)
    df = df[df.statuscode==200].reset_index(drop=True)
    if len(df)!=0:
        df_snaptimes = df[df.baseKey==df.baseKey4Yr].reset_index(drop=True)
        df_snaptimes = df_snaptimes.sort_values('datetime').groupby(['year']).tail(1).reset_index(drop=True)
        df_snaptimes = df_snaptimes[['year','date','datetime']]
        df_snaptimes.rename(columns={'date':'snapshot_date','datetime':'snapshot_datetime'}, inplace=True)
        if len(df_snaptimes)!=0:
            df_final = pd.DataFrame()
            for i in range(len(df_snaptimes)):
                df_int = df[df.date<=df_snaptimes.loc[i,'snapshot_date']]
                df_int = df_int[df_int.baseKey==df_int.baseKey4Yr]
                df_int = df_int.sort_values('date').groupby('url_id').tail(1).reset_index(drop=True)
                df_int['snapshot_date'] = df_snaptimes.loc[i,'snapshot_date']
                df_int['snapshot_datetime'] = df_snaptimes.loc[i,'snapshot_datetime']
                df_int['snapshot_year'] = pd.to_datetime(df_int['snapshot_date'], errors='coerce').dt.year
                df_int.rename(columns={'date':'url_date', 'year':'url_year', 'datetime':'url_datetime'}, inplace=True)
                df_int.drop(columns=['statuscode', 'baseKey', 'baseKey4Yr'], inplace=True)
                df_final = pd.concat([df_final,df_int], axis=0, ignore_index=True)
            df_final = df_final[['fm_id', 'snapshot_date', 'snapshot_datetime', 'snapshot_year', 'urlkey', 'url_date', 'url_datetime', 'url_year', 'url_id', 'url', 'length']]
            df_unique = df_final[['fm_id', 'url_datetime','url_id', 'url']]
            df_unique = df_unique.drop_duplicates(ignore_index=True)
            return df_final, df_unique
        else:
            return(None, None)
    else:
        return(None, None)

#PE Firms
cntr = 1
for file in os.listdir('./URLlist/PE Firms'):
    fileName = os.fsdecode(file).rstrip('.pkl')
    if os.path.isfile('./snapshotURLlevelList/PE Firms/%s.pkl' %fileName)==True:
        continue
    else:
        print('PE Firm', cntr, fileName)
        cntr = cntr+1
        df = pd.read_pickle('./URLlist/PE Firms/%s.pkl' %fileName, compression='zip')
        snapshotURLs, uniqueURLs = uniqueURLgen(df)
        if snapshotURLs is None:
            continue
        else:
            snapshotURLs.to_pickle('./snapshotURLlevelList/PE Firms/%s.pkl' %fileName, compression='zip')
            uniqueURLs.to_pickle('./scrapeList/PE Firms/%s.pkl' %fileName, compression='zip')
    
##HF Firms
cntr = 1
for file in os.listdir('./URLlist/HF Firms'):
    fileName = os.fsdecode(file).rstrip('.pkl')
    if os.path.isfile('./snapshotURLlevelList/HF Firms/%s.pkl' %fileName)==True:
        continue
    else:
        print('HF Firm', cntr, fileName)
        cntr = cntr+1
        df = pd.read_pickle('./URLlist/HF Firms/%s.pkl' %fileName, compression='zip')
        snapshotURLs, uniqueURLs = uniqueURLgen(df)
        if snapshotURLs is None:
            continue
        else:
            snapshotURLs.to_pickle('./snapshotURLlevelList/HF Firms/%s.pkl' %fileName, compression='zip')
            uniqueURLs.to_pickle('./scrapeList/HF Firms/%s.pkl' %fileName, compression='zip')