# %% Imports
import pandas as pd
import numpy as np
import os
import sys

sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
from jgtpy import JGTPDSP as pds
import tlid


from jgtutils.jgtos import get_data_path


# %% Functions


def _crop_dataframe(df, crop_last_dt: str = None, crop_start_dt: str = None):
    if crop_last_dt is not None:
        df = df[df.index <= crop_last_dt]
    if crop_start_dt is not None:
        df = df[df.index >= crop_start_dt]
    return df


def calculate_target_variable_min_max(
    dfsrc,
    crop_last_dt=None,
    crop_start_dt=None,
    WINDOW_MIN=1,
    WINDOW_MAX=150,
    set_index=True,
    rounder=2,
    pipsize=-1,
):
    df = dfsrc.copy()

    df = _crop_dataframe(df, crop_last_dt, crop_start_dt)

    # reset index before we iterate
    try:
        df.reset_index(drop=False, inplace=True)
    except:
        pass

    # Initialize the tmax and tmin columns with NaN values
    df["tmax"] = float("nan")
    df["tmin"] = float("nan")
    df["p"] = float("nan")
    df["l"] = float("nan")
    df["target"] = float("nan")

    # Calculate the maximum and minimum Close value in the window range for each row
    for i in range(WINDOW_MIN, len(df) - WINDOW_MAX):
        # FDBS
        if df.loc[i, "fdbs"] == 1.0:
            df.loc[i, "tmax"] = df.loc[i : i + WINDOW_MAX, "Close"].min()
            df.loc[i, "tmin"] = df.loc[i : i + WINDOW_MAX, "Close"].max()

            if df.loc[i, "High"] < df.loc[i, "tmin"]:
                df.loc[i, "l"] = round(df.loc[i, "High"] - df.loc[i, "Low"], rounder)
                df.loc[i, "target"] = round(
                    -1 * (df.loc[i, "High"] - df.loc[i, "Low"]), rounder
                )
            else:
                df.loc[i, "p"] = round(df.loc[i, "Low"] - df.loc[i, "tmax"], rounder)
                df.loc[i, "target"] = round(
                    df.loc[i, "Low"] - df.loc[i, "tmax"], rounder
                )
        # FDBB
        if df.loc[i, "fdbb"] == 1.0:
            df.loc[i, "tmax"] = df.loc[i : i + WINDOW_MAX, "Close"].max()
            df.loc[i, "tmin"] = df.loc[i : i + WINDOW_MAX, "Close"].min()

            if df.loc[i, "Low"] > df.loc[i, "tmin"]:
                df.loc[i, "l"] = round(df.loc[i, "High"] - df.loc[i, "Low"], rounder)
                df.loc[i, "target"] = round(
                    -1 * (df.loc[i, "High"] - df.loc[i, "Low"]), rounder
                )
            else:
                df.loc[i, "p"] = round(df.loc[i, "tmax"] - df.loc[i, "High"], rounder)
                df.loc[i, "target"] = round(
                    df.loc[i, "tmax"] - df.loc[i, "High"], rounder
                )

    # Fill NaN with zero for columns tmax and tmin
    df["tmax"] = df["tmax"].fillna(0)
    df["tmin"] = df["tmin"].fillna(0)
    df["p"] = df["p"].fillna(0)
    df["l"] = df["l"].fillna(0)
    df["target"] = df["target"].fillna(0)
    # After calculating the 'target' column
    if pipsize != -1:
        df["target"] = df["target"] / pipsize

    # @q Maybe set backnthe index !??
    if set_index:
        df.set_index("Date", inplace=True)
    return df


def pto_target_calculation(
    i, t, crop_start_dt=None, crop_end_dt=None, tlid_tag=None, output_report_dir=None,
    pto_vec_fdb_ao_vector_window_flag=False,   
    drop_calc_col=False,
    calc_col_to_drop_names = ["tmax", "tmin", "p", "l"],
    sel_1_suffix = "_sel",
    sel_1_keeping_columns = ["Low", "fdbs", "fdbb", "tmax", "tmin", "p", "l", "target"],
    sel_2_suffix = "_tnd",
    sel_2_keeping_columns = ["Open", "High", "Low", "Close", "fdbs", "fdbb", "target"],
    pto_vec_fdb_ao_out_s_name = 'vector_ao_fdbs',
    pto_vec_fdb_ao_out_b_name = 'vector_ao_fdbb',
    pto_vec_fdb_ao_in_s_sig_name = 'fdbs',
    pto_vec_fdb_ao_in_s_win_end_sig_name = 'zlcb',            
    pto_vec_fdb_ao_in_b_sig_name = 'fdbb',
    pto_vec_fdb_ao_in_b_win_end_sig_name = 'zlcs',
    pto_vec_fdb_ao_in_t_val_name = 'ao'
):
    """
    Calculate the PTO target based on the given POV parameters and output to file with report.

    Args:
        i (int): The value of i.get_fdb_ao_vector_window
        t (int): The value of t.
        crop_start_dt (datetime, optional): The start date for cropping. Defaults to None.
        crop_end_dt (datetime, optional): The end date for cropping. Defaults to None.
        tlid_tag (str, optional): The TLID tag. Defaults to None.
        output_report_dir (str, optional): The output report directory. Defaults to None.
        process_fdb_ao_vector_window (bool, optional): If True, process the fdb_ao_vector_window. Defaults to False.
        drop_calc_col (bool, optional): If True, drop the calculated columns. Defaults to False.
        calc_col_to_drop_names (list, optional): The list of column names to drop. Defaults to ["tmax", "tmin", "p", "l"].
        sel_1_suffix (str, optional): The suffix for the first selection. Defaults to "_sel".
        sel_1_keeping_columns (list, optional): The list of columns to keep for the first selection. Defaults to ["Low", "fdbs", "fdbb", "tmax", "tmin", "p", "l", "target"].
        sel_2_suffix (str, optional): The suffix for the second selection. Defaults to "_tnd".
        sel_2_keeping_columns (list, optional): The list of columns to keep for the second selection. Defaults to ["Open", "High", "Low", "Close", "fdbs", "fdbb", "target"].

    Returns:
        df_result_tmx (pandas.DataFrame): The DataFrame containing the calculated PTO target.
        sel1 (pandas.DataFrame): The DataFrame containing the selected columns.
        sel2 (pandas.DataFrame): The DataFrame containing the selected columns.
    """
    if tlid_tag is None:
        tlid_tag = tlid.get_minutes()

    default_jgtpy_data_full = "full/data"
    default_jgtpy_data_full = "/var/lib/jgt/full/data"
    data_dir_full = os.getenv("JGTPY_DATA_FULL", default_jgtpy_data_full)
    indir_cds = os.path.join(data_dir_full, "cds")
    outdir_tmx = os.path.join(
        data_dir_full, "targets", "mx"
    )  # @STCIssue Hardcoded path future JGTPY_DATA_FULL/.../mx
    df_result_tmx,sel1,sel2 =_pov_target_calculation_n_output240223(
        indir_cds,
        outdir_tmx,
        crop_start_dt,
        crop_end_dt,
        i,
        t,
        tlid_tag,
        output_report_dir=output_report_dir,
        pto_vec_fdb_ao_vector_window_flag=pto_vec_fdb_ao_vector_window_flag,
        drop_calc_col=drop_calc_col,
        calc_col_to_drop_names =calc_col_to_drop_names,
        sel_1_suffix = sel_1_suffix,
        sel_1_keeping_columns = sel_1_keeping_columns,
        sel_2_suffix = sel_2_suffix,
        sel_2_keeping_columns = sel_2_keeping_columns,
        pto_vec_fdb_ao_out_s_name = pto_vec_fdb_ao_out_s_name, 
        pto_vec_fdb_ao_out_b_name = pto_vec_fdb_ao_out_b_name,
        pto_vec_fdb_ao_in_s_sig_name = pto_vec_fdb_ao_in_s_sig_name,
        pto_vec_fdb_ao_in_s_win_end_sig_name = pto_vec_fdb_ao_in_s_win_end_sig_name,            
        pto_vec_fdb_ao_in_b_sig_name = pto_vec_fdb_ao_in_b_sig_name,
        pto_vec_fdb_ao_in_b_win_end_sig_name =pto_vec_fdb_ao_in_b_win_end_sig_name ,
        pto_vec_fdb_ao_in_t_val_name = pto_vec_fdb_ao_in_t_val_name
        
    )
    return df_result_tmx,sel1,sel2


def _pov_target_calculation_n_output240223(
    indir_cds,
    outdir_tmx,
    crop_start_dt,
    crop_end_dt,
    i,
    t,
    tlid_tag,
    output_report_dir=None,
    pto_vec_fdb_ao_vector_window_flag=False,
    drop_calc_col=True,
    write_reporting=True,
    calc_col_to_drop_names = ["tmax", "tmin", "p", "l"],
    sel_1_suffix = "_sel",
    sel_1_keeping_columns = ["High","Low", "fdbs", "fdbb", "tmax", "tmin", "p", "l", "target"],
    sel_2_suffix = "_tnd",
    sel_2_keeping_columns = ["Open", "High", "Low", "Close", "fdbs", "fdbb", "target"],
    pto_vec_fdb_ao_out_s_name = 'vector_ao_fdbs',
    pto_vec_fdb_ao_out_b_name = 'vector_ao_fdbb',
    pto_vec_fdb_ao_in_s_sig_name = 'fdbs',
    pto_vec_fdb_ao_in_s_win_end_sig_name = 'zlcb',            
    pto_vec_fdb_ao_in_b_sig_name = 'fdbb',
    pto_vec_fdb_ao_in_b_win_end_sig_name = 'zlcs',
    pto_vec_fdb_ao_in_t_val_name = 'ao'
):
    if tlid_tag is None:
        tlid_tag = tlid.get_minutes()

    ifn = i.replace("/", "-")
    # read instrument prop
    iprop = pds.get_instrument_properties(i)
    pipsize = iprop["pipsize"]
    nb_decimal = (
        len(str(pipsize).split(".")[1]) if "." in str(pipsize) else len(str(pipsize))
    )
    rounder = nb_decimal

    # Possible crop of the dataframe to a specific date range

    # Read the source data of already calculated CDS
    df_cds_source = pd.read_csv(
        f"{indir_cds}/{ifn}_{t}.csv", index_col=0, parse_dates=True
    )

    df_result_tmx = calculate_target_variable_min_max(
        df_cds_source, crop_end_dt, crop_start_dt, rounder=rounder, pipsize=pipsize
    )
        
    if pto_vec_fdb_ao_vector_window_flag:
        df_result_tmx = get_fdb_ao_vector_window(
            df_result_tmx,                                             
            out_s_name = pto_vec_fdb_ao_out_s_name, 
            out_b_name = pto_vec_fdb_ao_out_b_name,
            in_s_sig_name = pto_vec_fdb_ao_in_s_sig_name,
            in_s_win_end_sig_name = pto_vec_fdb_ao_in_s_win_end_sig_name,            
            in_b_sig_name = pto_vec_fdb_ao_in_b_sig_name,
            in_b_win_end_sig_name =pto_vec_fdb_ao_in_b_win_end_sig_name ,
            in_t_val_name = pto_vec_fdb_ao_in_t_val_name    
            )
        
        sel_1_keeping_columns.append(pto_vec_fdb_ao_out_s_name)
        sel_1_keeping_columns.append(pto_vec_fdb_ao_out_b_name)        
        # sel_1_keeping_columns.append([pto_vec_fdb_ao_out_s_name, pto_vec_fdb_ao_out_b_name,pto_vec_fdb_ao_in_s_sig_name, pto_vec_fdb_ao_in_s_win_end_sig_name, pto_vec_fdb_ao_in_b_sig_name, pto_vec_fdb_ao_in_b_win_end_sig_name, pto_vec_fdb_ao_in_t_val_name])        
        
        
        sel_2_keeping_columns.append(pto_vec_fdb_ao_out_s_name)
        sel_2_keeping_columns.append(pto_vec_fdb_ao_out_b_name)
        # sel_2_keeping_columns.append([pto_vec_fdb_ao_out_s_name, pto_vec_fdb_ao_out_b_name,pto_vec_fdb_ao_in_s_sig_name, pto_vec_fdb_ao_in_s_win_end_sig_name, pto_vec_fdb_ao_in_b_sig_name, pto_vec_fdb_ao_in_b_win_end_sig_name, pto_vec_fdb_ao_in_t_val_name])
        
        
        print(sel_1_keeping_columns)
        
        print(sel_2_keeping_columns)
        
    
    """  #@STCIssue We will Want those columns in the output selections
        out_s_name = 'vector_ao_fdbs',
        out_b_name = 'vector_ao_fdbb',
        in_s_sig_name = 'fdbs',
        in_s_win_end_sig_name = 'zlcb',            
        in_b_sig_name = 'fdbb',
        in_b_win_end_sig_name = 'zlcs',
        in_t_val_name = 'ao'
    """
    
    sel1 = df_result_tmx[sel_1_keeping_columns].copy()
    sel1 = sel1[(sel1["p"] != 0) | (sel1["l"] != 0)]
    
    output_sel_cols_fn = f"{outdir_tmx}/{ifn}_{t}{sel_1_suffix}.csv"
    try:
        sel1.to_csv(output_sel_cols_fn, index=True)
        print(f"Saved to {output_sel_cols_fn}")
    except Exception as e:
        print(f"Error occurred while saving to {output_sel_cols_fn}: {str(e)}")

    
    sel2 = df_result_tmx[sel_2_keeping_columns].copy()
    sel2["target"] = sel2["target"].round(rounder)
    sel2 = sel2[(sel2["target"] != 0)]
    
    output_tnd_targetNdata_fn = f"{outdir_tmx}/{ifn}_{t}{sel_2_suffix}.csv"
    try:
        sel2.to_csv(output_tnd_targetNdata_fn, index=True)
        print(f"Saved to {output_tnd_targetNdata_fn}")
    except Exception as e:
        print(f"Error occurred while saving to {output_tnd_targetNdata_fn}: {str(e)}")

    
    
    if drop_calc_col:
        df_result_tmx.drop(columns=calc_col_to_drop_names,inplace=True)
    # Save the result to a csv file
    output_all_cols_fn = f"{outdir_tmx}/{ifn}_{t}.csv"
    
    try:
        df_result_tmx.to_csv(output_all_cols_fn, index=True)
        print(f"Saved to {output_all_cols_fn}")
    except Exception as e:
        print(f"Error occurred while saving to {output_all_cols_fn}: {str(e)}")
    
    if write_reporting:
        _reporting(
        sel2, ifn, t, pipsize, tlid_tag, output_report_dir=output_report_dir
    )
        
    return df_result_tmx,sel1,sel2


def _reporting(df_selection2, ifn, t, pipsize, tlid_tag, output_report_dir=None):

    if output_report_dir is None:
        output_report_dir = os.getenv("JGTPY_DATA_FULL")

    report_file = f"{output_report_dir}/report-calc-{tlid_tag}.txt"
    print("Reporting to:", report_file)
    print(" tail -f ", report_file)

    with open(report_file, "a") as f:
        f.write(f"--- {ifn}_{t} --pipsize:{pipsize}---\n")
        f.write(f"Sum of target: {df_selection2['target'].sum()}\n\n")
        # f.write(f" (rounded): {(df_selection2['target'].sum()).round(2)}\n\n")


def readMXFile(
    instrument,
    timeframe,
    columns_to_remove=None,
    quiet=True,
    use_full=True,
    dt_crop_last=None,
    quote_count=None,
    sel_1_suffix = "_sel",
    sel_2_suffix = "_tnd",
    also_read_selections = False,
):
    """
    Read a MX Target file and return a pandas DataFrame.

    Parameters:
    instrument (str): The instrument name.
    timeframe (str): The timeframe of the data.
    columns_to_remove (list, optional): List of column names to remove from the DataFrame. Default is None.
    quiet (bool, optional): If True, suppresses the output messages. Default is True.
    use_full (bool, optional): If True, reads the full MX file. Default is True (there wont be MX in current data I think)).
    dt_crop_last (str, optional): The last date to crop the data. Default is None.
    quote_count (int, optional): The number of quotes to keep. Default is None.
    sel_1_suffix (str, optional): The suffix for the first selection. Defaults to "_sel".
    sel_2_suffix (str, optional): The suffix for the second selection. Defaults to "_tnd".
    also_read_selections (bool, optional): If True, also read the selections. Defaults to False.

    Returns:
    pandas.DataFrame: The DataFrame containing the MX Target data.
    OR
    tuple: A tuple containing the DataFrame and the selections DataFrames.
    """
    # Define the file path based on the environment variable or local path
    data_path_cds = get_data_path("targets/mx", use_full=use_full)
    fpath = pds.mk_fullpath(instrument, timeframe, "csv", data_path_cds)
    mdf = pd.read_csv(fpath)

    # Set 'Date' as the index and convert it to datetime
    mdf["Date"] = pd.to_datetime(mdf["Date"])
    mdf.set_index("Date", inplace=True)
    # Remove the specified columns
    if columns_to_remove is not None:
        mdf = mdf.drop(columns=columns_to_remove, errors="ignore")

    if dt_crop_last is not None:
        mdf = mdf[mdf.index < dt_crop_last]
    if quote_count is not None:
        mdf = mdf[-quote_count:]
    
    if also_read_selections:
        sel1 = pd.read_csv(fpath.replace(".csv", f"{sel_1_suffix}.csv"), index_col=0, parse_dates=True)
        sel2 = pd.read_csv(fpath.replace(".csv", f"{sel_2_suffix}.csv"), index_col=0, parse_dates=True)
        return mdf, sel1, sel2
    return mdf





# Upgrade to return a dataframe with the windows added as new columns to the original dataframe
def get_fdb_ao_vector_window(df,
                             out_s_name = 'vector_ao_fdbs',
                             out_b_name = 'vector_ao_fdbb',
                             in_s_sig_name = 'fdbs',
                             in_s_win_end_sig_name = 'zlcb',            
                             in_b_sig_name = 'fdbb',
                             in_b_win_end_sig_name = 'zlcs',
                             in_t_val_name = 'ao'
                             ):
    
    # reset index before we iterate
    try:
        df.reset_index(drop=False, inplace=True)
    except:
        pass
    
    
    df[out_s_name] = np.nan
    df[out_b_name] = np.nan
    for index, row in df.iterrows():
        if row[in_s_sig_name] == 1:
            window_start = index
            window_end = None
            for i in range(index, -1, -1):
                if df.at[i, in_s_win_end_sig_name] == 1:
                    window_end = i
                    break
            window = df.loc[window_end:window_start, in_t_val_name] if window_end is not None else df.loc[:window_start, in_t_val_name]
            df.at[index, out_s_name] = str(window.astype(float).tolist()) 
        if row[in_b_sig_name] == 1:
            window_start = index
            window_end = None
            for i in range(index, -1, -1):
                if df.at[i, in_b_win_end_sig_name] == 1:
                    window_end = i
                    break
            window = df.loc[window_end:window_start, in_t_val_name] if window_end is not None else df.loc[:window_start, in_t_val_name]
            df.at[index,out_b_name] = str(window.astype(float).tolist())
    #restore index
    try:
        df.set_index("Date", inplace=True)
    except:
        pass
    return df

