from typing import Dict, Sequence
import pandas as pd
from .utils import get_fdr, get_sd_quantile


def get_result_table(
    target_info_df: pd.DataFrame,
    param_hist_dict: Dict,
    guide_index: Sequence = None,
    suffix_fitted: str = "",
):
    """Return fitted result as table

    Args
        suffix_fitted: suffix the fitted parameters
    """
    if "quantiles" in param_hist_dict.keys():
        mu = param_hist_dict["quantiles"]["mu_alleles"][50].detach()[:, 0].numpy()
        mu_sd = get_sd_quantile(param_hist_dict["quantiles"]["mu_alleles"].detach())[
            :, 0
        ].numpy()
        sd = param_hist_dict["quantiles"]["sd_alleles"][50].detach()[:, 0].numpy()
    else:
        mu = param_hist_dict["params"]["mu_loc"].detach()[:, 0].numpy()
        mu_sd = param_hist_dict["params"]["mu_scale"].detach()[:, 0].numpy()
        sd = param_hist_dict["params"]["sd_loc"].detach().exp()[:, 0].numpy()
    mu_z = mu / mu_sd
    fdr_dec, fdr_inc, fdr, p_dec, p_inc, p = get_fdr(mu_z, get_p=True)
    fit_df = pd.DataFrame(
        {
            f"mu{suffix_fitted}": mu,
            f"mu_sd{suffix_fitted}": mu_sd,
            f"mu_z{suffix_fitted}": mu / mu_sd,
            f"sd{suffix_fitted}": sd,
            f"p_dec{suffix_fitted}": p_dec,
            f"p_inc{suffix_fitted}": p_inc,
            f"p{suffix_fitted}": p,
            f"fdr_dec{suffix_fitted}": fdr_dec,
            f"fdr_inc{suffix_fitted}": fdr_inc,
            f"fdr{suffix_fitted}": fdr,
        }
    )
    if "negctrl" in param_hist_dict.keys():
        print("Normalizing with common negative control distribution")
        mu0 = param_hist_dict["negctrl"]["mu_loc"]
        sd0 = param_hist_dict["negctrl"]["sd_loc"]
        print(f"Fitted mu0={mu0:3f}, sd0={sd0:3f}.")
        fit_df[f"mu_adj{suffix_fitted}"] = (mu - mu0) / sd0
        fit_df[f"mu_sd_adj{suffix_fitted}"] = mu_sd / sd0
        fit_df[f"sd_adj{suffix_fitted}"] = sd / sd0
    fit_df = pd.concat([target_info_df, fit_df], axis=1)

    if not "alpha_pi" in param_hist_dict["params"].keys():
        print(
            "Model didn't have fitted editing rate. Not writing the sgRNA result file."
        )
        sgRNA_df = None
    else:
        a_fitted = param_hist_dict["params"]["alpha_pi"].detach().numpy()
        pi = a_fitted[:, 0] / a_fitted.sum(axis=1)
        sgRNA_df = pd.DataFrame({f"edit_eff{suffix_fitted}": pi}, index=guide_index)
    return (fit_df, sgRNA_df)


def write_result_table(
    target_info_df, param_hist_dict, prefix="", guide_index=None, suffix_fitted=""
):
    """Write fitted result as table

    Args
        suffix_fitted: suffix the fitted parameters in the resulting table
    """
    fit_df, sgRNA_df = get_result_table(
        target_info_df,
        param_hist_dict,
        prefix=prefix,
        guide_index=guide_index,
        suffix_fitted=suffix_fitted,
    )
    fit_df.to_csv(f"{prefix}_variant_result.csv")
    if sgRNA_df:
        sgRNA_df.to_csv(f"{prefix}_CRISPRbean_sgRNA_result.csv")
    else:
        print("No sgRNA info fitted.")
