from dataclasses import dataclass
from typing import List
import matplotlib.pyplot as plt
import pandas as pd


@dataclass
class ChannelsGroup:
    channels: list
    color: str
    name: str = ""


@dataclass
class PlotData:
    channels_groups: List[ChannelsGroup]
    title: str
    x_axis_col_name: str
    y_axis_col_name: str
    alt_dashed: bool = False


cmap_to_color = {
    "Blues": "blue",
    "Greens": "green",
    "YlOrBr": "orange"
}


def merge_channels(df, channels_dict):
    new_df = df[df['ch'].isin(list(channels_dict.keys()))].copy()

    new_df['c_initial_mmol'] = new_df['ch'].map(
        lambda ch: channels_dict[ch][1])
    new_df['c_remaining_mmol'] = new_df['c_initial_mmol'] * \
        new_df['c_remaining_prop']

    new_df = new_df.groupby(['time_days']).agg(
        {
            'c_respired_prop': 'mean',
            'c_remaining_prop': 'mean',
            'c_initial_mmol': 'mean',
            'c_remaining_mmol': 'mean'
        }
    ).reset_index()

    return new_df


def aggregate_by_day(df):
    # For day 0, take the first value, for other days, take the mean
    day_0_df = df[df['time_days'] == 0].sort_values(
        by='time (days)').groupby(['ch', 'time_days']).first()
    other_days_df = df[df['time_days'] > 0].groupby(
        ['ch', 'time_days'])['c_respired_prop'].mean().reset_index()
    other_days_df = other_days_df.merge(df[['ch', 'time_days', 'int']].drop_duplicates([
                                        'ch', 'time_days']), on=['ch', 'time_days'])

    # Combine day 0 and other days
    df_one_row_per_day = pd.concat(
        [day_0_df.reset_index(), other_days_df], ignore_index=True)
    df_one_row_per_day = df_one_row_per_day.sort_values(
        by=['ch', 'time_days']).reset_index(drop=True)

    return df_one_row_per_day


def plot_c_prop_against_time(channel_groups, data_frame, substrate):
    fig = plt.figure()
    plt.figure(figsize=(8, 6))

    for cg in channel_groups:
        colors = [cmap_to_color[cg.color]] * len(cg.channels)
        styles = [':', '--', '-']
        for i, ch in enumerate(cg.channels):
            df_channel = data_frame[data_frame["ch"] == ch]
            (plt.plot(df_channel["time_days"], df_channel["c_respired_prop"],
                      label=f'CH{ch}', color=colors[i], linestyle=styles[i], linewidth=1.5))

    plt.rcParams.update({'mathtext.default':  'regular'})
    plt.title(
        f"Proportion of biodegraded C against time in days for all channels ({substrate})")
    plt.xlabel("Time (days)", fontsize=14)
    plt.xlim(0)
    plt.ylabel("Proportion of C respired", fontsize=14)
    plt.ylim(0, 1.0)
    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    plt.grid(color='grey', linestyle='--', linewidth=0.5)
    plt.tight_layout()
    plt.savefig(
        f'output/svg_plots/proprtion_biodegraded_all_channels_{substrate}.svg', format='svg')


def plot_merged_c_prop_against_time(data_frames, colors, substrate, temperatures):
    fig = plt.figure()
    plt.figure(figsize=(8, 6))

    for df_temp, temp, color in zip(data_frames, temperatures, colors):
        plt.scatter(df_temp['time_days'], df_temp['c_respired_prop'],
                    label=f'{temp}°C', color=color, s=7)
        plt.plot(df_temp['time_days'], df_temp['c_respired_prop'],
                 color=color, linewidth=1)

    plt.title(
        f'Proportion of biodegraded C against time in days for different temperatures ({substrate})')
    plt.xlabel('Time (days)', fontsize=14)
    plt.xlim(0)
    plt.ylabel('Proportion of C respired', fontsize=14)
    plt.ylim(0, 1.0)
    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    plt.grid(color='grey', linestyle='--', linewidth=0.5)
    plt.tight_layout()
    plt.savefig(
        f'output/svg_plots/proprtion_biodegraded_merged_channels_{substrate}.svg', format='svg')


def plot_c_remaining_mmol_against_time(data_frames, colors, temperatures, substrate):
    fig = plt.figure()
    plt.figure(figsize=(8, 6))

    for df, temp, color in zip(data_frames, temperatures, colors):
        plt.scatter(df['time_days'], df['c_remaining_mmol'],
                    label=f'{temp}°C', color=color, s=7)
        plt.plot(df['time_days'], df['c_remaining_mmol'],
                 color=color, linewidth=1)

    plt.title(
        f'Carbon remaining in mmoles versus time in days ({substrate})')
    plt.xlabel('Time (days)', fontsize=14)
    plt.xlim(0)
    plt.ylabel('Carbon remaining (mmoles)', fontsize=14)
    plt.ylim(0, 1.0)
    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    plt.grid(color='grey', linestyle='--', linewidth=0.5)
    plt.tight_layout()
    plt.savefig(
        f'output/svg_plots/c_remaining_{substrate}.svg', format='svg')


def plot_experimental_vs_fitted(df_all_t, color_map, regression, substrate):
    plt.figure(figsize=(10, 6))
    for T in df_all_t['temp'].unique():
        temp_df = df_all_t[df_all_t['temp'] == T]
        time_data = temp_df['time_days'].values
        c_c0_data = temp_df['c_remaining_prop'].values
        color = color_map.get(T)

        plt.scatter(time_data, c_c0_data, color=color,
                    label=f'Experimental Data at {T-273.15} °C')
        plt.plot(
            time_data, regression.make_predictions(time_data, T),
            linestyle='--', color=color, label=f'Fitted Curve at {T-273.15} °C')

    plt.xlabel('Time (days)')
    plt.xlim(0)
    plt.ylabel('Proportion of C remaining')
    plt.ylim(0, 1.0)
    plt.legend()
    plt.title(f'Global Kinetic Rate Law Fit for All Temperatures ({substrate})')
    plt.grid(color='grey', linestyle='--', linewidth=0.5)
    plt.tight_layout()
    plt.savefig(
        f'output/svg_plots/global_kinetic_rate_law_{substrate}.svg', format='svg')

    plt.show()


def get_predicted_values_to_print(temperatures, predicted_values, lag_days):
    # Specify tresholds
    thresholds = [
        ("50 %", 0.5),
        ("60 %", 0.4),
        ("70 %", 0.3),
        ("80 %", 0.2),
        ("90 %", 0.1),
        ("95 %", 0.05),
        ("99 %", 0.01)
    ]

    column_names = []
    for T in temperatures:
        column_names.append(f"Predicted № of Days ({T-273.15}°C)")
        column_names.append(f"Adjusted № of Days ({T-273.15}°C)")

    values_to_print = []
    for label, threshold in thresholds:
        values_for_threshold = [threshold, label]
        for T in temperatures:
            day_predicted = [i for i, n in enumerate(
                predicted_values[T]) if n <= threshold] 
            if day_predicted:
                values_for_threshold.append(day_predicted[0] + 1)
                values_for_threshold.append(day_predicted[0] + 1 + lag_days[T])
            else:
                values_for_threshold.append(None)
                values_for_threshold.append(None)
        values_to_print.append(values_for_threshold)

    df_values_to_print = pd.DataFrame(values_to_print,
                                      columns=["C/Co", "% of C biodegraded"] + column_names)

    df_values_to_print["C/Co"] = df_values_to_print["C/Co"].map(
        "{:.2f}".format)
    
    return df_values_to_print


def plot_predicted_values_vs_time(temperatures, labels, max_days, color_map,
                                  predicted_values, time_days, substrate):
    plt.figure(figsize=(10, 6))
    for T, label in zip(temperatures, labels):
        color = color_map.get(T)
        c_c0_values = predicted_values[T]

        # Only keep points that fall within the max_days limit after shifting
        valid_indices = [i for i, t in enumerate(time_days) if t <= max_days]

        if not valid_indices:  # Skip if no valid points
            continue

        # Filter both time and values using valid indices
        valid_time = [time_days[i] for i in valid_indices]
        valid_values = [c_c0_values[i] for i in valid_indices]

        # shifted_time = [t + days_to_add for t in time_days]
        plt.plot(valid_time, valid_values, color=color,
                 label=f'Predicted {label}', linestyle='--')

    plt.xlabel('Time (days)')
    plt.xlim(0)
    plt.ylabel(r"$C/C_0$")
    plt.ylim(0, 1.0)
    plt.legend()
    plt.title(
        f'Predicted C remaining versus time at various temperatures ({substrate})')
    plt.grid()

    plt.tight_layout()
    plt.savefig(
        f'output/svg_plots/predicted_time_{substrate}_{max_days}_days.svg', format='svg')

    plt.show()
