import pandas as pd
from typing import List
from utils import TEMP_DIR
import plotly.express as px
import plotly.io as pio
import os
from functions.chart_functions import scatter_chart_fig, llm_chart_data_scrub, _write_chart
from dotenv import load_dotenv

load_dotenv()

root_url = os.getenv("ROOT_URL", "")

def descriptive_stats_func(session_hash, session_folder, columns: List[str]=[], **kwargs):
    print("DESCRIPTIVE STATISTICS")
    try:
        from html import escape

        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
        csv_query_path = f'{dir_path}/query.csv'

        df = pd.read_csv(csv_query_path)

        if columns:
            df = df[[c for c in columns if c in df.columns]]

        desc = df.describe().round(4)

        header_cells = '<th style="background:#1e40af;">Statistic</th>' + ''.join(
            f'<th>{escape(str(col))}</th>' for col in desc.columns
        )
        row_html = [
            '<tr>'
            + f'<td style="font-weight:600;color:#1e40af;background:#eff6ff;white-space:nowrap;">{escape(str(idx))}</td>'
            + ''.join(f'<td>{escape(str(val))}</td>' for val in row)
            + '</tr>'
            for idx, row in desc.iterrows()
        ]

        style = (
            '<style>'
            '.vda-table-wrap{overflow-x:auto;margin:8px 0;border-radius:8px;border:1px solid #e5e7eb;}'
            '.vda-table{width:100%;border-collapse:collapse;font-size:13px;font-family:Inter,system-ui,sans-serif;}'
            '.vda-table thead th{background:#3B82F6;color:#fff;padding:9px 14px;text-align:left;white-space:nowrap;font-weight:600;}'
            '.vda-table tbody td{padding:7px 14px;border-bottom:1px solid #f1f5f9;white-space:nowrap;}'
            '.vda-table tbody tr:nth-child(even){background:#f8fafc;}'
            '.vda-table tbody tr:last-child td{border-bottom:none;}'
            '</style>'
        )
        table = (
            '<div class="vda-table-wrap"><table class="vda-table">'
            f'<thead><tr>{header_cells}</tr></thead>'
            '<tbody>' + '\n'.join(row_html) + '</tbody>'
            '</table></div>'
        )

        return {"reply": style + table}

    except Exception as e:
        print("DESCRIPTIVE STATS ERROR")
        print(e)
        return {"reply": f"There was an error generating descriptive statistics. Error: {e}. You should probably try again."}


def kmeans_clustering_func(feature_columns: List[str], x_column: str, y_column: str,
                           session_hash, session_folder, n_clusters: int = 3,
                           layout: List[dict] = [{}], **kwargs):
    print("KMEANS CLUSTERING")
    try:
        from sklearn.cluster import KMeans
        from sklearn.preprocessing import StandardScaler
        from html import escape

        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
        chart_path = f'{dir_path}/chart.html'
        csv_query_path = f'{dir_path}/query.csv'

        df = pd.read_csv(csv_query_path)

        feature_df = df[feature_columns].select_dtypes(include='number').dropna()
        if feature_df.shape[1] < 1:
            return {"reply": "No numeric feature columns found for clustering. Please refine your query to include numeric columns."}

        X_scaled = StandardScaler().fit_transform(feature_df)
        labels = KMeans(n_clusters=n_clusters, random_state=42, n_init=10).fit_predict(X_scaled)

        df_clustered = df.loc[feature_df.index].copy()
        df_clustered['Cluster'] = [f'Cluster {l}' for l in labels]

        fig = px.scatter(
            df_clustered, x=x_column, y=y_column, color='Cluster',
            title=f'K-Means Clustering (k={n_clusters})',
        )
        fig.update_layout(font=dict(family='Inter, system-ui, sans-serif'))

        _, layout_dict = llm_chart_data_scrub({}, layout)
        if layout_dict:
            fig.update_layout(**layout_dict)

        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
        iframe = _write_chart(fig, chart_path, chart_url)

        cluster_summary = df_clustered.groupby('Cluster')[feature_columns].mean().round(3)
        header_cells = '<th style="background:#1e40af;">Cluster</th>' + ''.join(
            f'<th>{escape(str(col))}</th>' for col in cluster_summary.columns
        )
        row_html = [
            '<tr>'
            + f'<td style="font-weight:600;color:#1e40af;background:#eff6ff;white-space:nowrap;">{escape(str(idx))}</td>'
            + ''.join(f'<td>{escape(str(val))}</td>' for val in row)
            + '</tr>'
            for idx, row in cluster_summary.iterrows()
        ]
        style = (
            '<style>'
            '.vda-table-wrap{overflow-x:auto;margin:8px 0;border-radius:8px;border:1px solid #e5e7eb;}'
            '.vda-table{width:100%;border-collapse:collapse;font-size:13px;font-family:Inter,system-ui,sans-serif;}'
            '.vda-table thead th{background:#3B82F6;color:#fff;padding:9px 14px;text-align:left;white-space:nowrap;font-weight:600;}'
            '.vda-table tbody td{padding:7px 14px;border-bottom:1px solid #f1f5f9;white-space:nowrap;}'
            '.vda-table tbody tr:nth-child(even){background:#f8fafc;}'
            '.vda-table tbody tr:last-child td{border-bottom:none;}'
            '</style>'
        )
        summary_table = (
            '<div class="vda-table-wrap"><table class="vda-table">'
            f'<thead><tr>{header_cells}</tr></thead>'
            '<tbody>' + '\n'.join(row_html) + '</tbody>'
            '</table></div>'
        )

        return {"reply": f'{iframe}\n\n**Cluster Centroids (feature means per cluster):**\n{style}{summary_table}'}

    except Exception as e:
        print("KMEANS CLUSTERING ERROR")
        print(e)
        return {"reply": f"There was an error running K-Means clustering. Error: {e}. You should probably try again."}


def hypothesis_test_func(test_type: str, column: str, session_hash, session_folder,
                         column2: str = "", group_column: str = "",
                         group_values: List[str] = [], pop_mean: float = 0.0, **kwargs):
    print("HYPOTHESIS TEST")
    try:
        from scipy import stats
        from html import escape

        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
        csv_query_path = f'{dir_path}/query.csv'
        df = pd.read_csv(csv_query_path)

        if test_type == "t_test_independent":
            if not group_column or group_column not in df.columns:
                return {"reply": "Please specify a valid group_column for the independent t-test."}
            unique_groups = df[group_column].dropna().unique().tolist()
            if group_values and len(group_values) == 2:
                g1_label, g2_label = group_values[0], group_values[1]
            elif len(unique_groups) == 2:
                g1_label, g2_label = unique_groups[0], unique_groups[1]
            else:
                return {"reply": f"For an independent t-test, exactly 2 groups are needed. Found: {unique_groups}. Specify group_values with 2 entries."}

            g1 = df[df[group_column] == g1_label][column].dropna()
            g2 = df[df[group_column] == g2_label][column].dropna()
            t_stat, p_value = stats.ttest_ind(g1, g2)

            result_rows = [
                ("Test", "Independent Samples T-Test"),
                ("Column", column),
                ("Group Column", group_column),
                (f"Group 1", str(g1_label)),
                (f"Group 2", str(g2_label)),
                (f"Group 1 Mean (n={len(g1)})", f"{g1.mean():.4f}"),
                (f"Group 2 Mean (n={len(g2)})", f"{g2.mean():.4f}"),
                ("T-Statistic", f"{t_stat:.4f}"),
                ("P-Value", f"{p_value:.6f}"),
                ("Significant at α=0.05", "Yes ✓" if p_value < 0.05 else "No ✗"),
            ]
            title = f"T-Test: {column} by {group_column}"

        elif test_type == "t_test_one_sample":
            sample = df[column].dropna()
            t_stat, p_value = stats.ttest_1samp(sample, pop_mean)
            result_rows = [
                ("Test", "One-Sample T-Test"),
                ("Column", column),
                ("Hypothesized Mean (μ₀)", f"{pop_mean:.4f}"),
                (f"Sample Mean (n={len(sample)})", f"{sample.mean():.4f}"),
                ("Sample Std Dev", f"{sample.std():.4f}"),
                ("T-Statistic", f"{t_stat:.4f}"),
                ("P-Value", f"{p_value:.6f}"),
                ("Significant at α=0.05", "Yes ✓" if p_value < 0.05 else "No ✗"),
            ]
            title = f"One-Sample T-Test: {column} vs μ={pop_mean}"

        elif test_type == "chi_square":
            if not column2 or column2 not in df.columns:
                return {"reply": "Please specify a valid column2 for the chi-square test."}
            contingency = pd.crosstab(df[column], df[column2])
            chi2, p_value, dof, _ = stats.chi2_contingency(contingency)
            result_rows = [
                ("Test", "Chi-Square Test of Independence"),
                ("Column 1", column),
                ("Column 2", column2),
                ("Chi-Square Statistic", f"{chi2:.4f}"),
                ("Degrees of Freedom", str(dof)),
                ("P-Value", f"{p_value:.6f}"),
                ("Significant at α=0.05", "Yes ✓" if p_value < 0.05 else "No ✗"),
            ]
            title = f"Chi-Square: {column} × {column2}"

        else:
            return {"reply": f"Unknown test_type '{test_type}'. Use one of: t_test_independent, t_test_one_sample, chi_square."}

        style = (
            '<style>'
            '.vda-table-wrap{overflow-x:auto;margin:8px 0;border-radius:8px;border:1px solid #e5e7eb;}'
            '.vda-table{width:100%;border-collapse:collapse;font-size:13px;font-family:Inter,system-ui,sans-serif;}'
            '.vda-table thead th{background:#3B82F6;color:#fff;padding:9px 14px;text-align:left;white-space:nowrap;font-weight:600;}'
            '.vda-table tbody td{padding:7px 14px;border-bottom:1px solid #f1f5f9;white-space:nowrap;}'
            '.vda-table tbody tr:nth-child(even){background:#f8fafc;}'
            '.vda-table tbody tr:last-child td{border-bottom:none;}'
            '</style>'
        )
        header_cells = f'<th style="background:#1e40af;" colspan="2">{escape(title)}</th>'
        row_html = [
            '<tr>'
            + f'<td style="font-weight:600;color:#1e40af;background:#eff6ff;white-space:nowrap;">{escape(label)}</td>'
            + f'<td>{escape(value)}</td>'
            + '</tr>'
            for label, value in result_rows
        ]
        table = (
            '<div class="vda-table-wrap"><table class="vda-table">'
            f'<thead><tr>{header_cells}</tr></thead>'
            '<tbody>' + '\n'.join(row_html) + '</tbody>'
            '</table></div>'
        )
        return {"reply": style + table}

    except Exception as e:
        print("HYPOTHESIS TEST ERROR")
        print(e)
        return {"reply": f"There was an error running the hypothesis test. Error: {e}. You should probably try again."}


def regression_func(independent_variables: List[str], dependent_variable: str, session_hash, session_folder, category: str='', **kwargs):
    print("LINEAR REGRESSION CALCULATION")
    print(independent_variables)
    print(dependent_variable)
    try:
        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
        chart_path = f'{dir_path}/chart.html'
        csv_query_path = f'{dir_path}/query.csv'

        df = pd.read_csv(csv_query_path)

        if category in df.columns:
           fig = scatter_chart_fig(df=df, x_column=independent_variables,y_column=dependent_variable, 
                                    category=category,trendline="ols")
        else:
           fig = scatter_chart_fig(df=df,x_column=independent_variables,y_column=dependent_variable, 
                                    trendline="ols")

        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
        iframe = _write_chart(fig, chart_path, chart_url)

        results_frame = px.get_trendline_results(fig)

        print("RESULTS")
        print(results_frame)
        print(results_frame.at[0, 'px_fit_results'])
        results = results_frame.at[0, 'px_fit_results']
        print(results.summary())

        return {"reply": '{"regression_summary": %s, "regression_chart": %s' % (str(results.summary()), str(iframe))} 
    
    except Exception as e:
      print("LINEAR REGRESSION ERROR")
      print(e)
      reply = f"""There was an error generating the linear regression calculation from {independent_variables} and {dependent_variable}
              The error is {e},
              You should probably try again.
              """
      return {"reply": reply}