# mcp_tools.py

import numpy as np
import pingouin as pg

import matplotlib.pyplot as plt
import io
import base64
import json
def plot_group_distribution(groups_str):
    groups = [np.array([float(x) for x in g.split(",")]) for g in groups_str.split(";")]

    fig, axes = plt.subplots(1, 3, figsize=(12, 4))

    # Histogram
    for g in groups:
        axes[0].hist(g, alpha=0.5)
    axes[0].set_title("Histogram")

    # Boxplot
    axes[1].boxplot(groups)
    axes[1].set_title("Boxplot")

    # QQ plot (first group)
    pg.qqplot(groups[0], ax=axes[2])
    axes[2].set_title("QQ Plot (Group 1)")

    buf = io.BytesIO()
    fig.savefig(buf, format="png")
    buf.seek(0)
    return base64.b64encode(buf.read()).decode()

def assumption_checks(groups_str):
    import pandas as pd

    groups = [np.array([float(x) for x in g.split(",")]) for g in groups_str.split(";")]


    # Normality (Shapiro) with safety guard
    normality = []
    
    for i, g in enumerate(groups):
        if len(g) < 4:
            normality.append({
                "group": f"G{i+1}",
                "error": "Normality test requires at least 4 samples."
            })
        else:
            result = pg.normality(g).to_dict(orient="records")[0]
            normality.append({
                "group": f"G{i+1}",
                **result
            })
    

    # Homogeneity (Levene)
    levene = pg.homoscedasticity(groups, method="levene").to_dict(orient="records")[0]

    return {
        "normality": normality,
        "homogeneity": levene
    }

def posthoc_tests(groups_str, method="games-howell"):
    import pandas as pd

    groups = [np.array([float(x) for x in g.split(",")]) for g in groups_str.split(";")]

    data = []
    for i, g in enumerate(groups):
        for val in g:
            data.append({"value": val, "group": f"G{i+1}"})

    df = pd.DataFrame(data)

    result = pg.pairwise_ttests(
        dv="value",
        between="group",
        data=df,
        padjust="bonf",
        parametric=(method == "tukey"),
        effsize="hedges"
    )

    return result.to_dict(orient="records")

# -----------------------------
# Helper: convert Pingouin DF → JSON
# -----------------------------
def pg_to_json(df):
    """Convert Pingouin single-row DataFrame to JSON dict."""
    return df.to_dict(orient="records")[0]


# -----------------------------
# Welch t-test (independent, unequal variances)
# -----------------------------
def welch_t_test(group1_str, group2_str, alternative="two-sided"):
    g1 = np.array([float(x) for x in group1_str.split(",")])
    g2 = np.array([float(x) for x in group2_str.split(",")])

    result = pg.ttest(g1, g2, paired=False, correction=True, alternative=alternative)
    return pg_to_json(result)


# -----------------------------
# Student t-test (independent, equal variances)
# -----------------------------
def student_t_test(group1_str, group2_str, alternative="two-sided"):
    g1 = np.array([float(x) for x in group1_str.split(",")])
    g2 = np.array([float(x) for x in group2_str.split(",")])

    result = pg.ttest(g1, g2, paired=False, correction=False, alternative=alternative)
    return pg_to_json(result)


# -----------------------------
# Paired t-test (same subjects measured twice)
# -----------------------------
def paired_t_test(group1_str, group2_str, alternative="two-sided"):
    g1 = np.array([float(x) for x in group1_str.split(",")])
    g2 = np.array([float(x) for x in group2_str.split(",")])

    result = pg.ttest(g1, g2, paired=True, alternative=alternative)
    return pg_to_json(result)


# -----------------------------
# One-sample t-test (sample vs population mean)
# -----------------------------
def one_sample_t_test(group_str, population_mean, alternative="two-sided"):
    g = np.array([float(x) for x in group_str.split(",")])

    # Adjust sample for one-sample t-test
    adjusted = g - float(population_mean)

    # Pingouin one-sample t-test: y=None, popmean implicitly = 0
    result = pg.ttest(adjusted, y=None, paired=False, alternative=alternative)

    return pg_to_json(result)


# -----------------------------
# One-way ANOVA (multiple independent groups)
# -----------------------------
def one_way_anova(groups_str):
    """
    groups_str format:
    "1,2,3; 4,5,6; 7,8,9"
    """
    groups = [np.array([float(x) for x in g.split(",")]) for g in groups_str.split(";")]

    # Build DataFrame for Pingouin
    import pandas as pd
    data = []
    for i, g in enumerate(groups):
        for val in g:
            data.append({"value": val, "group": f"G{i+1}"})

    df = pd.DataFrame(data)

    result = pg.anova(dv="value", between="group", data=df)
    return pg_to_json(result)