import argparse
from clearml import Task, Model
from src.policies.simple_baseline import BaselinePolicy, Battery
from src.data import DataProcessor, DataConfig
import torch
import numpy as np
import pandas as pd
import datetime
from tqdm import tqdm
from src.utils.imbalance_price_calculator import ImbalancePriceCalculator
import seaborn as sns
import matplotlib.pyplot as plt
import plotly.express as px

### import functions ###
from src.trainers.quantile_trainer import auto_regressive as quantile_auto_regressive
from src.trainers.diffusion_trainer import sample_diffusion
from src.utils.clearml import ClearMLHelper

### Arguments ###
parser = argparse.ArgumentParser()
parser.add_argument('--task_id', type=str, default=None)
parser.add_argument('--model_type', type=str, default=None)
parser.add_argument('--model_name', type=str, default=None)
args = parser.parse_args()

assert args.task_id is not None, "Please specify task id"
assert args.model_type is not None, "Please specify model type"
assert args.model_name is not None, "Please specify model name"

### Baseline Policy ###
battery = Battery(2, 1)
baseline_policy = BaselinePolicy(battery, data_path="")


### Load Imbalance Prices ###
imbalance_prices = pd.read_csv('data/imbalance_prices.csv', sep=';')
imbalance_prices["DateTime"] = pd.to_datetime(imbalance_prices['DateTime'], utc=True)
imbalance_prices = imbalance_prices.sort_values(by=['DateTime'])

def get_imbalance_prices(date):
    imbalance_prices_day = imbalance_prices[imbalance_prices["DateTime"].dt.date == date]
    return imbalance_prices_day['Positive imbalance price'].values

def load_model(task_id: str):
    """
    Load model from task id
    """
    task = Task.get_task(task_id=task_id)

    lstm = task.get_parameter("data_processor/lstm")
    full_day_skip = task.get_parameter("data_processor/full_day_skip")
    output_size = int(task.get_parameter("data_processor/output_size"))

    print(f"lstm: {lstm}")
    print(f"full_day_skip: {full_day_skip}")
    print(f"output_size: {output_size}")

    configuration = task.get_parameters_as_dict()
    data_features = configuration['data_features']

    ### Data Config ###
    data_config = DataConfig()
    for key, value in data_features.items():
        setattr(data_config, key, value == "True")
    print(data_config.__dict__)

    ### Data Processor ###
    data_processor = DataProcessor(data_config, path="", lstm=lstm=="True")
    data_processor.set_batch_size(8192)
    data_processor.set_full_day_skip(full_day_skip == "True")
    data_processor.set_output_size(int(output_size))

    ### Model ###
    output_model_id = task.output_models_id["checkpoint"]
    clearml_model = Model(model_id=output_model_id)
    filename = clearml_model.get_weights()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(filename)

    model.to(device)
    model.eval()

    _, test_loader = data_processor.get_dataloaders(
        predict_sequence_length=output_size
    )

    return configuration, model, data_processor, test_loader


def quantile_auto_regressive_predicted_NRV(model, date, data_processor, test_loader):
    idx = test_loader.dataset.get_idx_for_date(date.date())
    initial, _, samples, target = quantile_auto_regressive(test_loader.dataset, model, model.quantiles, [idx]*500, 96)
    samples = samples.cpu().numpy()
    target = target.cpu().numpy()

    # inverse using data_processor
    samples = data_processor.inverse_transform(samples)
    target = data_processor.inverse_transform(target)

    return initial.cpu().numpy()[0][-1], samples, target

def diffusion_predicted_NRV(model, date, _, test_loader):
    device = next(model.parameters()).device

    idx = test_loader.dataset.get_idx_for_date(date.date())
    prev_features, targets = test_loader.dataset.get_batch([idx])

    if len(list(prev_features.shape)) == 2:
        initial_sequence = prev_features[:, :96]
    else:
        initial_sequence = prev_features[:, :, 0]

    prev_features = prev_features.to(device)
    targets = targets.to(device)

    samples = sample_diffusion(model, 1000, prev_features)

    return initial_sequence.cpu().numpy()[0][-1], samples.cpu().numpy(), targets.cpu().numpy()

def get_next_day_profits_for_date(model, data_processor, test_loader, date, ipc, predict_NRV: callable, penalties: list):
    charge_thresholds = np.arange(-100, 250, 25)
    discharge_thresholds = np.arange(-100, 250, 25)

    predicted_nrv_profits_cycles = {i: [0, 0] for i in penalties}
    baseline_profits_cycles = {i: [0, 0] for i in penalties}

    _charge_thresholds = {}
    _discharge_thresholds = {}

    initial, nrvs, target = predict_NRV(model, date, data_processor, test_loader)

    initial = np.repeat(initial, nrvs.shape[0])
    combined = np.concatenate((initial.reshape(-1, 1), nrvs), axis=1)

    reconstructed_imbalance_prices = ipc.get_imbalance_prices_2023_for_date_vectorized(date, combined)
    reconstructed_imbalance_prices = torch.tensor(reconstructed_imbalance_prices, device="cuda")

    yesterday_imbalance_prices = get_imbalance_prices(date.date() - datetime.timedelta(days=1))
    yesterday_imbalance_prices = torch.tensor(np.array([yesterday_imbalance_prices]), device="cpu")

    real_imbalance_prices = get_imbalance_prices(date.date())

    for penalty in penalties:
        found_charge_thresholds, found_discharge_thresholds = baseline_policy.get_optimal_thresholds(reconstructed_imbalance_prices, charge_thresholds, discharge_thresholds, penalty)
        
        _charge_thresholds[penalty] = found_charge_thresholds
        _discharge_thresholds[penalty] = found_discharge_thresholds

        next_day_charge_threshold = found_charge_thresholds.mean(axis=0)
        next_day_discharge_threshold = found_discharge_thresholds.mean(axis=0)
        yesterday_charge_thresholds, yesterday_discharge_thresholds = baseline_policy.get_optimal_thresholds(yesterday_imbalance_prices, charge_thresholds, discharge_thresholds, penalty)


        next_day_profit, next_day_charge_cycles = baseline_policy.simulate(torch.tensor([[real_imbalance_prices]]), torch.tensor([next_day_charge_threshold]), torch.tensor([next_day_discharge_threshold]))
        yesterday_profit, yesterday_charge_cycles = baseline_policy.simulate(torch.tensor([[real_imbalance_prices]]), torch.tensor([yesterday_charge_thresholds.mean(axis=0)]), torch.tensor([yesterday_discharge_thresholds.mean(axis=0)]))

        predicted_nrv_profits_cycles[penalty][0] += next_day_profit.item()
        predicted_nrv_profits_cycles[penalty][1] += next_day_charge_cycles.item()

        baseline_profits_cycles[penalty][0] += yesterday_profit.item()
        baseline_profits_cycles[penalty][1] += yesterday_charge_cycles.item()

    return predicted_nrv_profits_cycles, baseline_profits_cycles, _charge_thresholds, _discharge_thresholds

def next_day_test_set(model, data_processor, test_loader, ipc, predict_NRV: callable):
    penalties = [0, 50, 250, 500, 1000, 1500]
    predicted_nrv_profits_cycles = {i: [0, 0] for i in penalties}
    baseline_profits_cycles = {i: [0, 0] for i in penalties}

    charge_thresholds = {}
    discharge_thresholds = {}

    dates = baseline_policy.test_data["DateTime"].dt.date.unique()
    dates = pd.to_datetime(dates)

    for date in tqdm(dates):
        try:
            new_predicted_nrv_profits_cycles, new_baseline_profits_cycles, new_charge_thresholds, new_discharge_thresholds = get_next_day_profits_for_date(model, data_processor, test_loader, date, ipc, predict_NRV, penalties)
            
            charge_thresholds[date] = new_charge_thresholds
            discharge_thresholds[date] = new_discharge_thresholds

            for penalty in penalties:
                predicted_nrv_profits_cycles[penalty][0] += new_predicted_nrv_profits_cycles[penalty][0]
                predicted_nrv_profits_cycles[penalty][1] += new_predicted_nrv_profits_cycles[penalty][1]

                baseline_profits_cycles[penalty][0] += new_baseline_profits_cycles[penalty][0]
                baseline_profits_cycles[penalty][1] += new_baseline_profits_cycles[penalty][1]

        except Exception as e:
            print(f"Error for date {date}")

    return predicted_nrv_profits_cycles, baseline_profits_cycles, charge_thresholds, discharge_thresholds

def main():
    clearml_helper = ClearMLHelper(project_name="Thesis/NrvForecast")
    task = clearml_helper.get_task(task_name="Policy Test")

    # task.execute_remotely(queue_name="default", exit_process=True)

    configuration, model, data_processor, test_loader = load_model(args.task_id)

    if args.model_type == "autoregressive_quantile":
        quantiles = configuration["general"]["quantiles"]
        quantiles = list(map(float, quantiles.strip('[]').split(',')))
        model.quantiles = quantiles
        predict_NRV = quantile_auto_regressive_predicted_NRV
        task.add_tags(["autoregressive_quantile"])
    elif args.model_type == "diffusion":
        predict_NRV = diffusion_predicted_NRV
        task.add_tags(["diffusion"])
    else:
        raise ValueError("Please specify model type")

    ipc = ImbalancePriceCalculator(data_path="")

    predicted_nrv_profits_cycles, baseline_profits_cycles, charge_thresholds, discharge_thresholds = next_day_test_set(model, data_processor, test_loader, ipc, predict_NRV)
    # the charge_thresholds is a dictionary with date as key. The values of the dictionary is another dictionary with keys as penalties and values as the charge thresholds
    # create density plot that shows a density plot of the charge thresholds for each penalty (use seaborn displot) (One plot with a different color for each penalty)

    charge_thresholds_for_penalty = {}
    for d in charge_thresholds.values():
        for penalty, thresholds in d.items():
            if penalty not in charge_thresholds_for_penalty:
                charge_thresholds_for_penalty[penalty] = []
            charge_thresholds_for_penalty[penalty].extend(thresholds)

    discharge_thresholds_for_penalty = {}
    for d in discharge_thresholds.values():
        for penalty, thresholds in d.items():
            if penalty not in discharge_thresholds_for_penalty:
                discharge_thresholds_for_penalty[penalty] = []
            discharge_thresholds_for_penalty[penalty].extend(thresholds)

    def plot_threshold_distribution(thresholds: dict, title: str):
        data_to_plot = []
        for penalty, values in thresholds.items():
            for value in values:
                data_to_plot.append({'Penalty': penalty, 'Value': value.item()})
        df = pd.DataFrame(data_to_plot)
        palette = sns.color_palette("bright", len(thresholds.keys()))
        fig = sns.displot(data=df, x="Value", hue="Penalty", kind="kde", palette=palette)
        plt.title('Density of Charge Thresholds by Penalty')
        plt.xlabel('Charge Threshold')
        plt.ylabel('Density')
        plt.legend(title='Penalty')
        task.get_logger().report_matplotlib_figure(
            "Policy Results", 
            title, 
            iteration=0, 
            figure=fig
        )
        plt.close()

    ### Plot charge thresholds distribution ###
    plot_threshold_distribution(charge_thresholds_for_penalty, "Charge Thresholds")

    ### Plot discharge thresholds distribution ###
    plot_threshold_distribution(discharge_thresholds_for_penalty, "Discharge Thresholds")

    def plot_thresholds_per_day(thresholds: dict, title: str):
        # plot mean charge threshold per day (per penalty (other color))
        data_to_plot = []
        for date, values in thresholds.items():
            for penalty, value in values.items():
                mean_val = value.mean().item()
                std_val = value.std().item()  # Calculate standard deviation
                data_to_plot.append({'Date': date, 'Penalty': penalty, 'Mean': mean_val, 'StdDev': std_val})
                print(f"Date: {date}, Penalty: {penalty}, Mean: {mean_val}, StdDev: {std_val}")
        df = pd.DataFrame(data_to_plot)
        df["Date"] = pd.to_datetime(df["Date"])

        fig = px.line(
            df,
            x="Date",
            y="Mean",
            color="Penalty",
            title=title,
            labels={"Mean": "Threshold", "Date": "Date"},
            markers=True,  # Adds markers to the lines
            hover_data=["Penalty"],  # Adds additional hover information
        )

        fig.update_layout(
            width=1000,  # Set the width of the figure
            height=600,  # Set the height of the figure
            title_x=0.5,  # Center the title horizontally
        )

        task.get_logger().report_plotly(
            "Thresholds per Day", 
            title, 
            iteration=0, 
            figure=fig
        )

    ### Plot mean charge thresholds per day ###
    plot_thresholds_per_day(charge_thresholds, "Mean Charge Thresholds per Day")

    ### Plot mean discharge thresholds per day ###
    plot_thresholds_per_day(discharge_thresholds, "Mean Discharge Thresholds per Day")


    # create dataframe with columns "name", "penalty", "profit", "cycles"
    df = pd.DataFrame(columns=["name", "penalty", "profit", "cycles"])

    # use concat
    for penalty in predicted_nrv_profits_cycles.keys():
        new_rows = pd.DataFrame({
            "name": [f"{args.model_type} ({args.model_name})", "baseline"],
            "penalty": [penalty, penalty],
            "profit": [predicted_nrv_profits_cycles[penalty][0], baseline_profits_cycles[penalty][0]],
            "cycles": [predicted_nrv_profits_cycles[penalty][1], baseline_profits_cycles[penalty][1]]
        })
        df = pd.concat([df, new_rows], ignore_index=True)

    # sort by name, penalty ascending
    df = df.sort_values(by=["name", "penalty"])

    # calculate profit per cycle
    df["profit_per_cycle"] = df.apply(lambda row: row["profit"] / row["cycles"] if row["cycles"] != 0 else 0, axis=1)

    task.get_logger().report_table(
        "Policy Results", 
        "Policy Results", 
        iteration=0, 
        table_plot=df
    )

    # plotly to show profit on y axis and cycles on x axis (show 2 lines, one for each model)
    fig = px.line(
        df, 
        x="cycles", 
        y="profit", 
        color="name", 
        title="Profit vs. Cycles for Each Model",
        labels={"cycles": "Cycles", "profit": "Profit"},
        markers=True,  # Adds markers to the lines
        hover_data=["penalty"]  # Adds additional hover information
    )

    fig.update_xaxes(autorange="reversed")


    task.get_logger().report_plotly(
        "Policy Results", 
        "Profit vs. Cycles for Each Model", 
        iteration=0, 
        figure=fig
    )

    # close task
    task.close()


if __name__ == "__main__":
    main()