"""A sciprt to drop optimization runs based on a run_date"""
import os
from ams.db import sql_execute

connection_string = os.environ['DB_A']

uk_models = [
    'UKGR',
    'UKCO',
    'UKCB',
    'UKBG',
    'UKBAL'
]

uk_models_policies = [
    'UKGR',
    'UKCO',
    'UKCB',
    'UKBG',
    'UKBAL',
    'PUKBAL',
    'PUKBG',
    'PUKCB',
    'PUKCO',
    'PUKGR'
]

uk_asset_classes = [
    'EMD_GBP',
    'UK_LARGE',
    'UK_MID',
    'UK_GOV',
    'US_AGG_GBP_H',
    'US_SHORT_GBP',
    'EUR_AGG_XUK_GBP_H',
    'PACIFIC_XJP_GBP',
    'EM_GBP',
    'US_HY_GBP_H',
    'US_LC_GBP',
    'UK_SHORT_GOV',
    'US_SHORT_GBP_H',
    'UK_CORP',
    'US_SMALL_GBP',
    'EUR_AGG_XUK_GBP',
    'EMD_GBP_H',
    'UK_CASH',
    'EUR_XUK_GBP',
    'JAPAN_GBP',
    'US_HY_GBP',
    'UK_LINK',
    'UK_SHORT_CORP',
    'US_AGG_GBP'
]

muni_models = [
    'HBGMN',
    'HBGMT',
    'HBMN',
    'HBMT',
    'HCBMN',
    'HCBMT',
    'HCMN',
    'HCMT'
]

muni_models_policies = [
    'HBGMN',
    'HBGMT',
    'HBMN',
    'HBMT',
    'HCBMN',
    'HCBMT',
    'HCMN',
    'HCMT',
    'PMUNBA',
    'PMUNBG',
    'PMUNCB',
    'PMUNCO'
]

muni_asset_classes = [
    'NT_MUNI',
    'MUNI_AA',
    'MUNI_BB',
    'MUNI_B',
    'MUNI_CCC',
    'EAFE_AC_MUNI',
    'US_LC_MUNI',
    'US_DEF_MUNI',
    'US_SMID_MUNI',
    'EAFE_MUNI',
    'EAFE_SC_MUNI',
    'EM_MUNI',
    'EMD_MUNI',
    'INTL_AGG_MUNI',
    'EMLC_MUNI',
    'CASH_MUNI',
    'MA_MUNI',
    'MF_MUNI',
    'US_LONG_MUNI',
    'US_MID_MUNI',
    'US_SMALL_MUNI',
    'MUNI',
    'MUNI_HY',
    'MUNI_SHORT',
    'MUNI_AAA',
    'MUNI_A',
    'MUNI_BBB'
]

uk_models = tuple(uk_models)
uk_models_policies = tuple(uk_models_policies)
uk_asset_classes = tuple(uk_asset_classes)
muni_models = tuple(muni_models)
muni_models_policies = tuple(muni_models_policies)
muni_asset_classes = tuple(muni_asset_classes)

def delete_fi_cmas(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(
        connection_string,
        f"""DELETE FROM ams_equilibrium_cma WHERE run_date = '{run_date}'"""
        )
    sql_execute(
        connection_string,
        f"""DELETE FROM ams_fi_cma WHERE run_date = '{run_date}'"""
        )
    sql_execute(
        connection_string,
        f"""DELETE FROM ams_fi_inputs WHERE run_date = '{run_date}'"""
        )


def delete_pre_opt(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_asset_class_stats WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_distributions WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_matrices WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_model_stats WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_scenarios WHERE run_date = '{run_date}'""")


def delete_opt(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_account WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_policy WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_rebal WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_rebal WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_account WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_policy WHERE run_date = '{run_date}'""")


def delete_opt_to(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_to_weights_rebal WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_to_stats_rebal WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_to_weights_acct WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_to_stats_acct WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_to_weights_pol WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_to_stats_pol WHERE run_date = '{run_date}'""")


def delete_pre_opt_eqinc(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_eq_inc_matrices WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_eq_inc_scenarios WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_eq_inc_fund_stats WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_eq_inc_coefficients WHERE run_date = '{run_date}'""")


def delete_opt_eqinc(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_eqinc_ax_weights_rebal WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_eqinc_ax_stats_rebal WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_eqinc_ax_weights_account WHERE run_date = '{run_date}'""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_eqinc_ax_stats_account WHERE run_date = '{run_date}'""")


def delete_opt_uk(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_account 
                WHERE run_date = '{run_date}' 
                AND model_code in 
                {repr(uk_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_policy 
                WHERE run_date = '{run_date}' 
                AND model_code in 
                {repr(uk_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_rebal 
                WHERE run_date = '{run_date}' 
                AND model_code in 
                {repr(uk_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_rebal 
                WHERE run_date = '{run_date}' 
                AND model_code in 
                {repr(uk_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_account 
                WHERE run_date = '{run_date}' 
                AND model_code in 
                {repr(uk_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_policy 
                WHERE run_date = '{run_date}' 
                AND model_code in 
                {repr(uk_models)}""")


def delete_pre_opt_uk(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_asset_class_stats 
                WHERE run_date = '{run_date}' 
                AND asset_class in 
                {repr(uk_asset_classes)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_distributions 
                WHERE run_date = '{run_date}' 
                AND asset_class in 
                {repr(uk_asset_classes)}""")
    sql_execute(connection_string,
    f"""DELETE FROM ams_matrices 
        WHERE run_date = '{run_date}' 
        AND (asset_class_1 in 
        {repr(uk_asset_classes)} 
        OR asset_class_2 in 
        {repr(uk_asset_classes)})""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_model_stats 
                WHERE run_date = '{run_date}' 
                AND model_code in 
                {repr(uk_models_policies)}""")
    sql_execute(connection_string,
        f"""DELETE FROM ams_scenarios 
        WHERE run_date = '{run_date}' 
        AND asset_class in 
        {repr(uk_asset_classes)}""")

def delete_opt_muni(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_account 
                WHERE run_date = '{run_date}' 
                AND model_code in {repr(muni_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_policy 
                WHERE run_date = '{run_date}' 
                AND model_code in {repr(muni_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_weights_rebal 
                WHERE run_date = '{run_date}' 
                AND model_code in {repr(muni_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_rebal W
                HERE run_date = '{run_date}' 
                AND model_code in {repr(muni_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_account 
                WHERE run_date = '{run_date}' 
                AND model_code in {repr(muni_models)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_axioma_stats_policy 
                WHERE run_date = '{run_date}' 
                AND model_code in {repr(muni_models)}""")


def delete_pre_opt_muni(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_asset_class_stats 
                WHERE run_date = '{run_date}' 
                AND asset_class in 
                {repr(muni_asset_classes)}""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_distributions 
                WHERE run_date = '{run_date}' 
                AND asset_class in 
                {repr(muni_asset_classes)}""")
    sql_execute(connection_string, f"""DELETE FROM ams_matrices 
        WHERE run_date = '{run_date}' 
        AND (asset_class_1 in 
        {repr(muni_asset_classes)}
        OR asset_class_2 in 
        {repr(muni_asset_classes)})""")
    sql_execute(connection_string,
                f"""DELETE FROM ams_model_stats 
                WHERE run_date = '{run_date}' 
                AND model_code in 
                {repr(muni_models_policies)}""")
    sql_execute(connection_string, f"""DELETE FROM ams_scenarios 
        WHERE run_date = '{run_date}' 
        AND asset_class in 
        {repr(muni_asset_classes)}""")


def delete_raw_returns(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_raw_returns WHERE time_series_date = '{run_date}'""")


def delete_fundamentals(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_fundamentals WHERE run_date = '{run_date}'""")


def delete_equilibrium(run_date):
    """
    This function deletes SQL output from the script in the function name for a given run_date

    Args:
    run_date (str): The date to delete from the SQL tables
    """
    sql_execute(connection_string,
                f"""DELETE FROM ams_equilibrium_cma WHERE run_date = '{run_date}'""")


def delete_optimization_output(script, run_date):
    """
    This function deletes SQL output from the optimization script runs for a given run_date.
    List of Scripts [delete_raw_returns, delete_fundamentals, delete_equilibrium, delete_fi_cmas, 
    delete_pre_opt, delete_opt, delete_opt_to, delete_pre_opt_eqinc, delete_opt_eqinc, delete_opt_uk,
    delete_pre_opt_uk, delete_opt_muni, delete_pre_opt_muni]

    Args:
    script (str): The script that was run to delete from the SQL tables
    run_date (str): The date to delete from the SQL tables
    """
    run_funcs = {
        'delete_raw_returns', 
        'delete_fundamentals', 
        'delete_equilibrium', 
        'delete_fi_cmas', 
        'delete_pre_opt', 
        'delete_opt', 
        'delete_opt_to', 
        'delete_pre_opt_eqinc', 
        'delete_opt_eqinc', 
        'delete_opt_uk',
        'delete_pre_opt_uk', 
        'delete_opt_muni', 
        'delete_pre_opt_muni'
        }
    assert script in run_funcs, 'script not in dictionary'
    func_return = globals()[script](run_date)
    return func_return


DELETE_DATE = '30-SEP-2023'
delete_optimization_output('delete_opt', DELETE_DATE)