import pandas as pd
import argparse
from datetime import datetime


def remove_otc(df):
    return df[df["Exchange"] != "OTC"]


def reorder_cols(df):
    return df[
        [
            "Company Name",
            "Ticker",
            "Exchange",
            "Sector",
            "Industry",
            "Month of Fiscal Yr End",
            "Market Cap (mil)",
            "Last Close",
            "F0 Consensus Est.",
            "F1 Consensus Est.",
            "F2 Consensus Est.",
        ]
    ]


def rename_cols(df):
    return df.rename(
        columns={
            "F0 Consensus Est.": "EPS0",
            "F1 Consensus Est.": "EPS1",
            "F2 Consensus Est.": "EPS2",
        }
    )

def add_cols(df):
    df['EG1'] = ((df['EPS1'] - df['EPS0']) / abs(df['EPS1'])).round(4)
    df['EG2'] = ((df['EPS2'] - df['EPS1']) / abs(df['EPS2'])).round(4)
    df['PE0'] = (df['Last Close'] / df['EPS0']).round(2)
    df['PE1'] = (df['Last Close'] / df['EPS1']).round(2)
    df['PE2'] = (df['Last Close'] / df['EPS2']).round(2)
    df['PEG1'] = (df['PE1'] / (df['EG1'] * 100)).round(2)
    df['PEG2'] = (df['PE2'] / (df['EG2'] * 100)).round(2)
    df['EG1 Mean'] = df.groupby('Industry')['EG1'].transform('mean').round(4)
    df['EG2 Mean'] = df.groupby('Industry')['EG2'].transform('mean').round(4)
    df['PE0 Mean'] = df.groupby('Industry')['PE0'].transform('mean').round(2)
    df['PE1 Mean'] = df.groupby('Industry')['PE1'].transform('mean').round(2)
    df['PE2 Mean'] = df.groupby('Industry')['PE2'].transform('mean').round(2)
    return df


if __name__ == "__main__":
    today = datetime.today()
    iso_date = today.strftime("%Y-%m-%d")
    parser = argparse.ArgumentParser(
        "zacks-screen", description="Create cleaned CSV file from zacks screen"
    )
    parser.add_argument(
        "filename",
        help="Path to zacks screen CSV file",
    )
    parser.add_argument(
        "--dest",
        default=f"zacks-screen-{iso_date}.csv",
        help="Location to save parsed CSV to",
    )
    args = parser.parse_args()
    df = pd.read_csv(args.filename)
    df = remove_otc(df)
    df = reorder_cols(df)
    df = rename_cols(df)
    df = add_cols(df)
    df.to_csv(args.dest, index=False)
    print(df)