import os
from datetime import datetime

import numpy as np
import pandas as pd
from glob import glob
import arakawa as ar
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import ttest_ind, chi2_contingency

from photonai import Hyperpipe
from photonai.processing import ResultsHandler

from photonai_projects.utils import find_latest_photonai_run


class Reporter:
    def __init__(self, project_folder):
        self.project_folder = project_folder
        self.md_file = os.path.join(self.project_folder, 'README.md')
        self.plot_dir = os.path.join(self.project_folder, 'plots')
        self.analysis_type = None
        self.best_config_metric = None

        if not os.path.exists(self.md_file):
            with open(self.md_file, 'w') as file:
                content = f"""# PHOTONAI Project {self.project_folder}

Add project description here            
"""
                file.write(content)

    def descriptive_sample_statistics(self, df: pd.DataFrame, title: str, footnote: str,
                                      categorical_target: str, continuous_variables: list,
                                      categorical_variables: list, lancet_format: bool = False):

        list_of_group_tables = list()
        groups = df[categorical_target].unique()
        for group_name in groups:
            df_group = df[df[categorical_target] == group_name]
            table = pd.DataFrame(columns=[group_name])

            # =====================
            # Categories
            # =====================
            for cat_variable in categorical_variables:
                categories = df_group[cat_variable].unique()
                row = pd.DataFrame(index=[cat_variable])
                table = pd.concat([table, row])
                for current_category in categories:
                    n_group = np.sum(df_group[cat_variable] == current_category)
                    n = df_group.shape[0]
                    content = "{} ({:.1%})".format(n_group, n_group / n)
                    row = pd.DataFrame({group_name: content}, index=[current_category])
                    table = pd.concat([table, row])

            # =====================
            # Continuous Variables
            # =====================
            for variable_name in continuous_variables:
                content = "{:.2f} ({:.2f})".format(df_group[variable_name].mean(), df_group[variable_name].std())
                row = pd.DataFrame({group_name: content}, index=[variable_name])
                table = pd.concat([table, row])

            list_of_group_tables.append(table)

        final_table = pd.concat(list_of_group_tables, axis=1)

        # =====================
        # Difference stat test
        # =====================

        diff_test_table = pd.DataFrame(columns=['Difference'])

        for cat_variable in categorical_variables:
            obs = pd.crosstab(index=df[cat_variable], columns=df[categorical_target])[groups]
            chi2, p, dof, ex = chi2_contingency(obs)
            content = self.format_p_value(p)
            row = pd.DataFrame({'Difference': content}, index=[cat_variable])
            diff_test_table = pd.concat([diff_test_table, row])

        for variable_name in continuous_variables:
            p = ttest_ind(a=df.loc[df[categorical_target] == groups[0], variable_name],
                          b=df.loc[df[categorical_target] == groups[1], variable_name], nan_policy='omit').pvalue
            content = self.format_p_value(p)
            row = pd.DataFrame({'Difference': content}, index=[variable_name])
            diff_test_table = pd.concat([diff_test_table, row])

        final_table = pd.concat([final_table, diff_test_table], axis=1)

        # ===========================
        # Handle decimal if necessary
        # ===========================
        if lancet_format:
            final_table = final_table.astype(str)
            for key in final_table.keys():
                final_table[key] = final_table[key].str.replace('.', '·')

        # =====================
        # Title and footnote
        # =====================
        if footnote:
            row = pd.DataFrame(index=[footnote])
            final_table = pd.concat([final_table, row])

        final_table = final_table.T.reset_index().T.reset_index()
        description_row = pd.DataFrame({'index': title}, index=[0])
        final_table = pd.concat([description_row, final_table])

        # last clean-ups
        final_table.iloc[1, 0] = ""
        final_table.iloc[0, 1:] = ""
        final_table.iloc[-1, 1:] = ""

        return final_table

    def format_p_value(self, p):
        if p < 0.001:
            content = "<0.001"
        elif p < 0.01:
            content = "{:.3f}".format(p)
        else:
            content = "{:.2f}".format(p)
        return content

    def save_sample_descriptives(self, df: pd.DataFrame,
                                 categorical_target: str,
                                 continuous_covariates: list,
                                 categorical_covariates: list):
        df.to_csv(os.path.join(self.project_folder, 'df_sample.csv'), index=False)

        table = self.descriptive_sample_statistics(df,
                                              title="Descriptive statistics",
                                              footnote="*t or χ² tests.",
                                              categorical_target=categorical_target,
                                              continuous_variables=continuous_covariates,
                                              categorical_variables=categorical_covariates,
                                              lancet_format=False)
        table.to_csv(os.path.join(self.project_folder, 'sample_description.csv'), index=False)

        # create plots describing the sample
        os.makedirs(self.plot_dir, exist_ok=True)
        plt.figure()
        target_dist = sns.countplot(data=df, x=categorical_target, alpha=0.75)
        plt.savefig(os.path.join(self.plot_dir, 'sample__target_distribution.png'))

        for cov in continuous_covariates:
            plt.figure()
            plot = sns.histplot(data=df, x=cov, hue=categorical_target,
                                kde=True, multiple='dodge')
            plt.savefig(os.path.join(self.plot_dir, f'sample__{cov}_distribution.png'))

        for cov in categorical_covariates:
            plt.figure()
            plot = sns.countplot(data=df, x=categorical_target, hue=cov, alpha=0.75)
            plt.savefig(os.path.join(self.plot_dir, f'sample__{cov}_distribution.png'))

        return table

    @staticmethod
    def metric_type(metric_name):
        """
        Check whether a metric name corresponds to a classification or regression metric.
        """
        classification_metrics = {
            "accuracy", "balanced_accuracy", "f1", "f1_macro", "f1_micro",
            "f1_weighted", "precision", "precision_macro", "precision_micro",
            "precision_weighted", "recall", "recall_macro", "recall_micro",
            "recall_weighted", "roc_auc", "roc_auc_ovr", "roc_auc_ovo",
            "log_loss", "matthews_corrcoef", "jaccard", "jaccard_score",
        }

        regression_metrics = {
            "r2", "neg_mean_absolute_error", "neg_mean_squared_error",
            "neg_root_mean_squared_error", "neg_median_absolute_error",
            "neg_mean_squared_log_error", "explained_variance",
            "max_error", "mean_absolute_error", "mean_squared_error",
            "median_absolute_error",
        }

        name = metric_name.lower()
        if name in classification_metrics:
            return "classification"
        elif name in regression_metrics:
            return "regression"
        else:
            return "unknown"

    def collect_results(self):
        results = list()
        analysis_folders = glob(os.path.join(self.project_folder, '*/'))
        analysis_folders = sorted(analysis_folders)
        for analysis in analysis_folders:
            if os.path.basename(os.path.dirname(analysis)) == 'plots':
                continue
            folder = find_latest_photonai_run(analysis)

            if folder is None:
                continue
            handler = ResultsHandler()
            handler.load_from_file(os.path.join(folder, "photonai_results.json"))
            pipe_infos = f"""
# PHOTONAI Analysis Info  
**Name**: {handler.results.name}  
**PHOTONAI Version**: {handler.results.version}

**No. of Samples**: {handler.results.hyperpipe_info.data['X_shape'][0]}  
**No. of Features**: {handler.results.hyperpipe_info.data['X_shape'][1]}

**Outer Cross-Validation**: {handler.results.hyperpipe_info.cross_validation['OuterCV']}  
**Inner Cross-Validation**: {handler.results.hyperpipe_info.cross_validation['InnerCV']}

**Hyperparameter Optimization**: {handler.results.hyperpipe_info.optimization['Optimizer']}  
**Primary Evaluation Metric**: {handler.results.hyperpipe_info.optimization['BestConfigMetric']}
 
"""
            with open(os.path.join(folder, "hyperpipe_infos.md"), 'w') as file:
                file.write(pipe_infos)

            # save important results and infos on hyperpipe
            df = handler.get_performance_table()
            df.to_csv(os.path.join(folder, "metrics.csv"), index=False)
            df = df.iloc[-1, :]
            best_config_metric = handler.results.hyperpipe_info.best_config_metric
            self.analysis_type = self.metric_type(best_config_metric)
            self.best_config_metric = best_config_metric
            best_metric = pd.DataFrame({'name': [best_config_metric],
                                       'value': [df[best_config_metric]]})
            best_metric.to_csv(os.path.join(folder, "best_metric.csv"))

            names = df.index.tolist()
            df['analysis'] = os.path.basename(os.path.dirname(analysis))
            df['photonai_folder'] = folder
            df = df[['analysis', 'photonai_folder'] + names]
            results.append(df)
        df = pd.DataFrame(results)
        df = df.drop(columns=['best_config', 'n_train', 'n_validation', 'fold'])
        df.to_csv(os.path.join(self.project_folder, 'summary.csv'), index=False, float_format='%.4f')

    def add_header(self, header: str, content: list):
        dp_header = ar.Text(f"## {header}")
        blocks = list()
        blocks.append(dp_header)
        blocks.extend(content)
        return ar.Group(blocks=blocks, columns=1)

    def write_report(self):
        df = pd.read_csv(os.path.join(self.project_folder, 'summary.csv'))

        project_page = ar.Group(
            label="Project",
            blocks=[ar.Group(ar.Text(
                    "<img src='https://avatars.githubusercontent.com/u/63720198?s=400&u=17e0e95ba5f7a7a220cfaaadce784094f8429478&v=4' alt='drawing' width='100'/>"),
                    #"![](https://avatars.githubusercontent.com/u/63720198?s=400&u=17e0e95ba5f7a7a220cfaaadce784094f8429478&v=4)"),
                    ar.Text(file=self.md_file),
                    columns=1)],
                 columns=1)

        descriptives_page = None
        if os.path.exists(os.path.join(self.project_folder, "df_sample.csv")):
            data_df = pd.read_csv(os.path.join(self.project_folder, 'df_sample.csv'))
            desc_df = pd.read_csv(os.path.join(self.project_folder, 'sample_description.csv'))
            descriptives_table = self.add_header(header="Descriptive Statistics",
                                                 content=[ar.DataTable(desc_df, label="Descriptive Statistics")])
            data_table = self.add_header(header="Sample Data",
                                         content=[ar.DataTable(data_df, label="Sample Data")])
            plot1 = self.add_header(header="target distribution", content=[ar.Media(file=os.path.join(self.plot_dir, 'sample__target_distribution.png'),
                             label="target distribution")])
            plots = glob(os.path.join(self.plot_dir, 'sample__*'))
            dp_plots = list()
            dp_plots.append(plot1)
            for plot in plots:
                if plot == os.path.join(self.plot_dir, 'sample__target_distribution.png'):
                    continue
                else:
                    dp_plots.append(self.add_header(header=f"{os.path.basename(plot).split('sample__')[-1].split('.png')[0].replace('_', ' ')}",
                                                    content=[ar.Media(file=plot,
                             label=f"{os.path.basename(plot)}")]))
            page_description = """
This page summarizes descriptive statistics of the sample used in the ML analyses. The first table presents
the phenotypic information on the sample. The second table summarizes descriptive statistics. The distribution
of the main descriptive variables are illustrated below (separately for the groups in the predictive target). 
            """
            descriptives_page = ar.Group(
                label="Sample Descriptives",
                blocks=[ar.Group(blocks=[ar.Text(page_description), data_table, descriptives_table], columns=1),
                    ar.Group(blocks=dp_plots, columns=2)])

        # results summary page
        df_short_results = df[['analysis', self.best_config_metric, f'{self.best_config_metric}_sem']]
        plt.figure(figsize=(10, 6))
        summary_plot = sns.barplot(
                    data=df_short_results, x=self.best_config_metric, y="analysis",
                   color="#6597B8")
        if self.analysis_type == 'classification':
            plt.xlim([0.3, 1])
        for index, row in df_short_results.iterrows():
            plt.errorbar(
                x=row[self.best_config_metric],
                y=index,
                xerr=row[f'{self.best_config_metric}_sem'],
                color='black',  # Match the point color
                capsize=0,  # Size of the error bar caps
                capthick=1,  # Thickness of the error bar caps
            )
        if self.analysis_type == 'classification':
            plt.axvline(0.5, color='grey', linestyle='--')
        plt.xlabel(f"{self.best_config_metric}")
        plt.ylabel("Analyses")

        summary_page = ar.Group(
            label="Results Summary",
            blocks=[
                self.add_header(header="Short Summary of PHOTONAI Analyses",
                                content=[ar.Plot(summary_plot,
                                                 label='Summary Plot',
                                                 caption="Mean performance across folds. Error bars depict +- 1 SD.")]),
                self.add_header(header="Detailed Summary of PHOTONAI Analyses",
                                content=[ar.DataTable(df, label="PHOTONAI Results", caption="Summary of PHOTONAI Analyses")])],
                 columns=1)

        # individual photonai analysis results
        pipelines = list()
        for _, analysis in df.iterrows():
            pipe_results = pd.read_csv(os.path.join(analysis['photonai_folder'], 'metrics.csv'))
            overall_pipe_results = pipe_results.iloc[-1, :]
            pipe_results.drop([pipe_results.iloc[-1].name], axis=0, inplace=True)
            best_metric = pd.read_csv(os.path.join(analysis['photonai_folder'], 'best_metric.csv'))

            overall_pipe_results = overall_pipe_results.drop(['best_config', 'fold', 'n_train', 'n_validation'])
            primary_metric_group = ar.Group(blocks=[ar.BigNumber(heading=f"{best_metric['name'][0]}", value=f"{best_metric['value'][0]:.2} [+-{overall_pipe_results[best_metric['name'][0] + '_sem']:.2}]")],
                                            columns=1)

            pipelines.append(ar.Group(blocks=[
                ar.Group(blocks=[
                ar.Text(file=os.path.join(analysis['photonai_folder'], "hyperpipe_infos.md")),
                self.add_header(header="Primary Metric (Mean [SD])",
                                content=[primary_metric_group])], columns=2),
                self.add_header(header="PHOTONAI Results across Folds",
                                content=[ar.DataTable(pipe_results)])], label=f"{analysis['analysis']}"))
        if len(pipelines) == 1:
            pipe_page = ar.Group(blocks=pipelines,
                                 label="Detailed Results")
        else:
            pipe_page = ar.Group(blocks=[ar.Text("Choose analysis..."),
                                         ar.Select(blocks=pipelines, type=ar.SelectType.DROPDOWN)],
                                 label="Detailed Results")

        #report = ar.View(ar.Select(blocks=[project_page, descriptives_page, summary_page, pipe_page]))
        report = ar.View(ar.Select(blocks=[project_page, summary_page, pipe_page]))
        ar.save_report(report, path=os.path.join(self.project_folder, "report.html"), open=True)
