import re

import nilearn
import numpy as np
import pandas as pd
from mne.utils import warn
from scipy import stats


def glm_to_tidy(info, statistic, design_matrix, wide=True, order=None):
    """
    Export GLM regression or contrast results in tidy format.

    Creates a long pandas data frame from regression results or contrast
    as computed by run_glm or compute_contrast.

    Parameters
    ----------
    info : MNE.Info
        Instance of MNE info.
    statistic : nilearn data,
        Either dict of nilearn.stats.regression.RegressionResults as returned
        by run_glm, or nilearn.stats.contrasts.Contrast as returned by
        compute_contrast.
    design_matrix : DataFrame
        As specified in Nilearn.
    wide : Bool
        Should the returned dataframe be in wide format. If False, then the
        returned data will be in long format.
    order : list
        Order that the channels should be returned with.

    Returns
    -------
    df : Tidy data frame,
        Data from statistic object in tidy data form.
    """
    if isinstance(statistic, dict) and isinstance(
        statistic[list(statistic.keys())[0]], nilearn.glm.regression.RegressionResults
    ):
        df = _tidy_RegressionResults(info, statistic, design_matrix)

    elif isinstance(statistic, nilearn.glm.contrasts.Contrast):
        df = _tidy_Contrast(info, statistic, design_matrix)

    else:
        raise TypeError(
            "Unknown statistic type. Expected dict of RegressionResults "
            f"or Contrast type. Received {type(statistic)}"
        )

    if wide:
        df = _tidy_long_to_wide(df, expand_output=True)

    if order is not None:
        df["old_index"] = df.index
        df = df.set_index("ch_name")
        df = df.loc[order, :]
        df["ch_name"] = df.index
        df.index = df["old_index"]
        df.drop(columns="old_index", inplace=True)
        df.rename_axis(None, inplace=True)

    return df


def _tidy_Contrast(data, glm_est, design_matrix):
    df = pd.DataFrame()
    if hasattr(glm_est, "stat_type"):
        contrast_type = glm_est.stat_type
    else:  # nilearn < 0.10
        contrast_type = glm_est.contrast_type
    effect = glm_est.effect
    if effect.ndim == 2:  # old nilearn
        assert effect.shape[0] == 1
        effect = effect[0]
    for idx, ch in enumerate(data.ch_names):
        df = pd.concat(
            [
                df,
                pd.DataFrame(
                    {
                        "ch_name": ch,
                        "ContrastType": contrast_type,
                        "variable": "effect",
                        "value": effect[idx],
                    },
                    index=[0],
                ),
                pd.DataFrame(
                    {
                        "ch_name": ch,
                        "ContrastType": contrast_type,
                        "variable": "p_value",
                        "value": glm_est.p_value()[idx],
                    },
                    index=[1],
                ),
                pd.DataFrame(
                    {
                        "ch_name": ch,
                        "ContrastType": contrast_type,
                        "variable": "stat",
                        "value": glm_est.stat()[idx],
                    },
                    index=[2],
                ),
                pd.DataFrame(
                    {
                        "ch_name": ch,
                        "ContrastType": contrast_type,
                        "variable": "z_score",
                        "value": glm_est.z_score()[idx],
                    },
                    index=[3],
                ),
            ],
            ignore_index=True,
        )
    return df


def _tidy_RegressionResults(data, glm_est, design_matrix):
    if not (data.ch_names == list(glm_est.keys())):
        warn("MNE data structure does not match regression results")

    theta_estimates = np.zeros((len(glm_est), len(design_matrix.columns)))
    t_estimates = np.zeros((len(glm_est), len(design_matrix.columns)))
    df_estimates = np.zeros((len(glm_est), len(design_matrix.columns)))
    p_estimates = np.zeros((len(glm_est), len(design_matrix.columns)))
    mse_estimates = np.zeros((len(glm_est), len(design_matrix.columns)))
    se_estimates = np.zeros((len(glm_est), len(design_matrix.columns)))

    for idx, name in enumerate(glm_est.keys()):
        theta_estimates[idx, :] = glm_est[name].theta.T
        df_estimates[idx, :] = glm_est[name].df_model
        mse_estimates[idx, :] = glm_est[name].MSE[0]
        for cond_idx, cond in enumerate(design_matrix.columns):
            t_estimates[idx, cond_idx] = glm_est[name].t(column=cond_idx).item()
            p_estimates[idx, cond_idx] = 2 * stats.t.cdf(
                -1.0 * np.abs(t_estimates[idx, cond_idx]),
                df=df_estimates[idx, cond_idx],
            )
            se_estimates[idx, cond_idx] = np.sqrt(np.diag(glm_est[name].vcov()))[
                cond_idx
            ]

    list_vals = [0] * (len(data.ch_names) * len(design_matrix.columns) * 6)
    idx = 0
    for ch_idx, ch in enumerate(data.ch_names):
        for cond_idx, cond in enumerate(design_matrix.columns):
            list_vals[0 + idx] = {
                "ch_name": ch,
                "Condition": cond,
                "variable": "theta",
                "value": theta_estimates[ch_idx][cond_idx],
            }
            list_vals[1 + idx] = {
                "ch_name": ch,
                "Condition": cond,
                "variable": "t",
                "value": t_estimates[ch_idx][cond_idx],
            }
            list_vals[2 + idx] = {
                "ch_name": ch,
                "Condition": cond,
                "variable": "df",
                "value": df_estimates[ch_idx][cond_idx],
            }
            list_vals[3 + idx] = {
                "ch_name": ch,
                "Condition": cond,
                "variable": "p_value",
                "value": p_estimates[ch_idx][cond_idx],
            }
            list_vals[4 + idx] = {
                "ch_name": ch,
                "Condition": cond,
                "variable": "mse",
                "value": mse_estimates[ch_idx][cond_idx],
            }
            list_vals[5 + idx] = {
                "ch_name": ch,
                "Condition": cond,
                "variable": "se",
                "value": se_estimates[ch_idx][cond_idx],
            }
            idx += 6

    dict_vals, i = {}, 0
    for entry in list_vals:
        dict_vals[i] = {
            "ch_name": entry["ch_name"],
            "Condition": entry["Condition"],
            "variable": entry["variable"],
            "value": entry["value"],
        }
        i = i + 1
    df = pd.DataFrame.from_dict(dict_vals, "index")

    return df


def _tidy_long_to_wide(d, expand_output=True):
    indices = ["ch_name"]
    if "Condition" in d.columns:
        # Regression results have a column condition
        indices.append("Condition")
    if "ContrastType" in d.columns:
        # Regression results have a column condition
        indices.append("ContrastType")

    d = d.set_index(indices)
    d = d.pivot_table(columns="variable", values="value", index=indices)
    d.reset_index(inplace=True)

    if expand_output:
        try:
            d["Source"] = [
                int(re.search(r"S(\d+)_D(\d+) (\w+)", ch).group(1))
                for ch in d["ch_name"]
            ]
            d["Detector"] = [
                int(re.search(r"S(\d+)_D(\d+) (\w+)", ch).group(2))
                for ch in d["ch_name"]
            ]
            d["Chroma"] = [
                re.search(r"S(\d+)_D(\d+) (\w+)", ch).group(3) for ch in d["ch_name"]
            ]
        except AttributeError:
            warn("Non standard source detector names used")
        d["Significant"] = d["p_value"] < 0.05

    return d
