"""Formatting for DR-DiD results."""

from typing import NamedTuple

import numpy as np
from scipy.stats import norm

from moderndid.core.format import (
    attach_format,
    compute_significance,
    format_footer,
    format_section_header,
    format_single_result_table,
    format_title,
)


def format_did_result(result: NamedTuple) -> str:
    """Format a DR-DiD result for display."""
    att = result.att
    se = result.se
    lci = result.lci
    uci = result.uci
    args = getattr(result, "args", {})

    result_type_name = type(result).__name__
    estimator_type = args.get("type", _infer_estimator_type(result_type_name))
    est_method = args.get("est_method", args.get("estMethod", "default"))
    panel = args.get("panel", _infer_panel_type(result_type_name))
    boot = args.get("boot", False)
    nboot = args.get("nboot", 999)
    boot_type = args.get("boot_type", "weighted")

    t_stat = att / se if se > 0 and not np.isnan(se) else np.nan
    p_value = 2 * (1 - norm.cdf(abs(t_stat))) if not np.isnan(t_stat) else np.nan

    lines = []

    title = _get_estimator_title(estimator_type, est_method)
    lines.extend(format_title(title))

    if hasattr(result, "call_params") and "data_shape" in result.call_params:
        n_obs, n_cols = result.call_params["data_shape"]
        lines.append(f" Computed from {n_obs} observations and {n_cols - 4} covariates.")

    sig_marker = compute_significance(lci, uci)

    lines.extend(
        format_single_result_table(
            "ATT",
            att,
            se,
            95,
            lci,
            uci,
            p_value=p_value,
            sig_marker=sig_marker,
        )
    )

    lines.extend(format_section_header("Data Info"))
    if panel is not None:
        data_type = "Panel data" if panel else "Repeated cross-sections"
        lines.append(f" Data structure: {data_type}")

    lines.extend(format_section_header("Estimation Details"))
    method_details = _get_method_description(estimator_type, est_method, args)
    for detail in method_details:
        lines.append(f" {detail}")

    lines.extend(format_section_header("Inference"))
    if boot:
        lines.append(f" Standard errors: Bootstrapped ({nboot} replications)")
        lines.append(f" Bootstrap type: {boot_type.capitalize()}")
    else:
        lines.append(" Standard errors: Analytical")

    if "trim_level" in args and estimator_type in ["ipw", "dr"]:
        lines.append(f" Propensity score trimming: {args['trim_level']}")

    if np.isnan(att) or (se is not None and np.isnan(se)):
        lines.append("")
        lines.append("Warning: Estimation failed. Check for data issues:")
        if "warning_msg" in args:
            lines.append(f"     {args['warning_msg']}")
        else:
            lines.append("     - Insufficient variation in treatment/control groups")
            lines.append("     - Perfect separation in propensity score")
            lines.append("     - Collinearity in covariates")

    reference = None
    if estimator_type == "dr":
        reference = "Reference: Sant'Anna and Zhao (2020), Journal of Econometrics"
    elif estimator_type == "ipw":
        reference = "Reference: Abadie (2005), Review of Economic Studies"
    elif estimator_type == "or":
        reference = "Reference: Heckman et al. (1997), Review of Economic Studies"

    lines.extend(format_footer(reference))

    return "\n".join(lines)


def _infer_estimator_type(result_type_name: str) -> str:
    name_lower = result_type_name.lower()
    if "drdid" in name_lower:
        return "dr"
    if "ipwdid" in name_lower or "stdipw" in name_lower:
        return "ipw"
    if "regdid" in name_lower or "ordid" in name_lower:
        return "or"
    if "twfe" in name_lower:
        return "twfe"
    return "unknown"


def _infer_panel_type(result_type_name: str) -> bool | None:
    name_lower = result_type_name.lower()
    if "panel" in name_lower:
        return True
    if "rc" in name_lower:
        return False
    return None


def _get_estimator_title(estimator_type: str, est_method: str) -> str:
    titles = {
        "dr": {
            "imp": "Doubly Robust DiD Estimator (Improved Method)",
            "trad": "Doubly Robust DiD Estimator (Traditional Method)",
            "imp_local": "Doubly Robust DiD Estimator (Improved Locally Efficient Method)",
            "trad_local": "Doubly Robust DiD Estimator (Traditional Local Method)",
            "default": "Doubly Robust DiD Estimator",
        },
        "ipw": {
            "ipw": "Inverse Probability Weighted DiD Estimator",
            "std_ipw": "Standardized IPW DiD Estimator (Hajek-type)",
            "default": "Inverse Probability Weighted DiD Estimator",
        },
        "or": {"default": "Outcome Regression DiD Estimator"},
        "twfe": {"default": "Two-Way Fixed Effects DiD Estimator"},
    }

    if estimator_type in titles:
        return titles[estimator_type].get(est_method, titles[estimator_type]["default"])
    return "Difference-in-Differences Estimator"


def _get_method_description(estimator_type: str, est_method: str, args: dict) -> list[str]:
    details = []

    if estimator_type == "dr":
        if est_method == "imp":
            details.append("Outcome regression: Weighted least squares")
            details.append("Propensity score: Inverse probability tilting")
        elif est_method == "trad":
            details.append("Outcome regression: Ordinary least squares")
            details.append("Propensity score: Logistic regression")
        elif est_method == "imp_local":
            details.append("Outcome regression: Locally weighted least squares")
            details.append("Propensity score: Inverse probability tilting")
        elif est_method == "trad_local":
            details.append("Outcome regression: Local linear regression")
            details.append("Propensity score: Logistic regression")
        else:
            details.append("Outcome regression and propensity score models")

    elif estimator_type == "ipw":
        normalized = args.get("normalized", True)
        if est_method == "std_ipw" or normalized:
            details.append("Weight type: Normalized (Hajek-type estimator)")
        else:
            details.append("Weight type: Unnormalized (Horvitz-Thompson-type)")
        details.append("Propensity score: Logistic regression")

    elif estimator_type == "or":
        details.append("Estimation method: Ordinary least squares")

    elif estimator_type == "twfe":
        details.append("Fixed effects: Unit and time")
        details.append("Estimation method: Within transformation")

    return details


def print_did_result(result_class):
    """Attach the DR-DiD formatter to a result class."""
    attach_format(result_class, format_did_result)
    return result_class
