"""Plotnine-based plotting functions for moderndid."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Literal

from plotnine import (
    aes,
    element_text,
    facet_wrap,
    geom_errorbar,
    geom_hline,
    geom_line,
    geom_point,
    geom_ribbon,
    geom_vline,
    ggplot,
    labs,
    position_dodge,
    scale_color_manual,
    scale_x_continuous,
    theme,
)

from moderndid.did.aggte_obj import AGGTEResult
from moderndid.did.multiperiod_obj import MPResult
from moderndid.didinter.results import DIDInterResult
from moderndid.didtriple.agg_ddd_obj import DDDAggResult
from moderndid.didtriple.estimators.ddd_mp import DDDMultiPeriodResult
from moderndid.didtriple.estimators.ddd_mp_rc import DDDMultiPeriodRCResult
from moderndid.plots.converters import (
    aggteresult_to_polars,
    dddaggresult_to_polars,
    dddmpresult_to_polars,
    didinterresult_to_polars,
    doseresult_to_polars,
    honestdid_to_polars,
    mpresult_to_polars,
    pteresult_to_polars,
)
from moderndid.plots.themes import COLORS, theme_minimal

if TYPE_CHECKING:
    from moderndid.didcont.estimation.container import DoseResult, PTEResult
    from moderndid.didhonest.honest_did import HonestDiDResult


def plot_gt(
    result: MPResult | DDDMultiPeriodResult | DDDMultiPeriodRCResult,
    show_ci: bool = True,
    ref_line: float | None = 0,
    title: str = "Group",
    xlab: str | None = None,
    ylab: str | None = None,
    ncol: int = 1,
    **_kwargs: Any,
) -> ggplot:
    """Plot group-time average treatment effects.

    Parameters
    ----------
    result : MPResult, DDDMultiPeriodResult, or DDDMultiPeriodRCResult
        Multi-period result object containing group-time ATT estimates.
        This should be the output from ``att_gt()`` or ``ddd()``.
    show_ci : bool, default=True
        Whether to show confidence intervals as error bars.
    ref_line : float or None, default=0
        Y-value for reference line. Set to None to hide.
    title : str, default="Group"
        Title prefix for each facet panel.
    xlab : str, optional
        X-axis label. Defaults to "Time".
    ylab : str, optional
        Y-axis label. Defaults to "ATT".
    ncol : int, default=1
        Number of columns in the facet grid. Use 1 for vertical stacking.

    Returns
    -------
    ggplot
        A plotnine ggplot object that can be further customized.
    """
    if isinstance(result, MPResult):
        df = mpresult_to_polars(result)
        plot_title = "Group-Time Average Treatment Effects"
    elif isinstance(result, (DDDMultiPeriodResult, DDDMultiPeriodRCResult)):
        df = dddmpresult_to_polars(result)
        plot_title = "Group-Time DDD Treatment Effects"
    else:
        raise TypeError(
            f"plot_gt requires MPResult, DDDMultiPeriodResult, or DDDMultiPeriodRCResult, got {type(result).__name__}"
        )

    df = df.with_columns([df["group"].cast(int).cast(str).alias("group_label")])
    x_breaks = sorted(df["time"].unique().to_list())

    plot = (
        ggplot(df, aes(x="time", y="att", color="treatment_status"))
        + geom_point(size=3, alpha=0.8)
        + scale_color_manual(
            values={"Pre": COLORS["pre_treatment"], "Post": COLORS["post_treatment"]},
            limits=["Pre", "Post"],
            name="Treatment Status",
        )
        + scale_x_continuous(breaks=x_breaks)
        + facet_wrap("~group_label", ncol=ncol, labeller=lambda x: f"{title} {x}", scales="free_x")
        + labs(
            x=xlab or "Time",
            y=ylab or "ATT",
            title=plot_title,
        )
        + theme_minimal()
        + theme(
            strip_text=element_text(size=11, weight="bold"),
            plot_title=element_text(margin={"b": 25}),
            legend_position="bottom",
        )
    )

    if show_ci:
        plot = plot + geom_errorbar(
            aes(ymin="ci_lower", ymax="ci_upper"),
            width=0.3,
            alpha=0.7,
        )

    if ref_line is not None:
        plot = plot + geom_hline(yintercept=ref_line, linetype="dashed", color="black", alpha=0.5)

    return plot


def plot_event_study(
    result: AGGTEResult | PTEResult | DDDAggResult,
    show_ci: bool = True,
    ref_line: float | None = 0,
    ref_period: float | None = None,
    xlab: str | None = None,
    ylab: str | None = None,
    title: str | None = None,
    **_kwargs: Any,
) -> ggplot:
    """Create event study plot for dynamic treatment effects.

    Parameters
    ----------
    result : AGGTEResult, PTEResult, or DDDAggResult
        Aggregated treatment effect result with dynamic/eventstudy aggregation,
        or PTEResult with event_study attribute.
    show_ci : bool, default=True
        Whether to show confidence intervals as error bars.
    ref_line : float or None, default=0
        Y-value for reference line. Set to None to hide.
    ref_period : float or None, default=-1
        X-value for vertical reference period line. Set to None to hide.
    xlab : str, optional
        X-axis label. Defaults to "Event Time".
    ylab : str, optional
        Y-axis label. Defaults to "ATT".
    title : str, optional
        Plot title. Defaults based on result type.

    Returns
    -------
    ggplot
        A plotnine ggplot object that can be further customized.
    """
    if hasattr(result, "event_study") and not isinstance(result, (AGGTEResult, DDDAggResult)):
        df = pteresult_to_polars(result)
        default_title = "Event Study"
    elif isinstance(result, DDDAggResult):
        if result.aggregation_type != "eventstudy":
            raise ValueError(f"Event study plot requires eventstudy aggregation, got {result.aggregation_type}")
        df = dddaggresult_to_polars(result)
        default_title = "DDD Event Study"
    elif isinstance(result, AGGTEResult):
        if result.aggregation_type != "dynamic":
            raise ValueError(f"Event study plot requires dynamic aggregation, got {result.aggregation_type}")
        df = aggteresult_to_polars(result)
        default_title = "Event Study"
    else:
        raise TypeError(
            f"plot_event_study requires AGGTEResult, PTEResult, or DDDAggResult, got {type(result).__name__}"
        )

    plot = ggplot(df, aes(x="event_time", y="att"))

    if show_ci:
        plot = plot + geom_errorbar(
            aes(ymin="ci_lower", ymax="ci_upper", color="treatment_status"),
            width=0.2,
            size=0.8,
        )

    if ref_line is not None:
        plot = plot + geom_hline(yintercept=ref_line, linetype="dashed", color="#7f8c8d", alpha=0.7)

    if ref_period is not None:
        plot = plot + geom_vline(xintercept=ref_period, linetype="dotted", color="#7f8c8d", alpha=0.7)
    else:
        plot = plot + geom_line(color=COLORS["line"], size=0.8, alpha=0.6, linetype="dotted")

    x_breaks = sorted(df["event_time"].unique().to_list())

    plot = (
        plot
        + geom_point(aes(color="treatment_status"), size=3.5)
        + scale_color_manual(
            values={"Pre": COLORS["pre_treatment"], "Post": COLORS["post_treatment"]},
            limits=["Pre", "Post"],
            name="Treatment Status",
        )
        + scale_x_continuous(breaks=x_breaks)
        + labs(
            x=xlab or "Event Time",
            y=ylab or "ATT",
            title=title or default_title,
        )
        + theme_minimal()
        + theme(legend_position="bottom")
    )

    return plot


def plot_agg(
    result: AGGTEResult | DDDAggResult,
    show_ci: bool = True,
    ref_line: float | None = 0,
    xlab: str | None = None,
    ylab: str | None = None,
    title: str | None = None,
    **_kwargs: Any,
) -> ggplot:
    """Create plot for aggregated treatment effects by group or calendar time.

    Parameters
    ----------
    result : AGGTEResult or DDDAggResult
        Aggregated treatment effect result with group or calendar aggregation.
    show_ci : bool, default=True
        Whether to show confidence intervals as error bars.
    ref_line : float or None, default=0
        Y-value for reference line. Set to None to hide.
    xlab : str, optional
        X-axis label. Defaults based on aggregation type.
    ylab : str, optional
        Y-axis label. Defaults to "ATT".
    title : str, optional
        Plot title. Defaults based on aggregation type.

    Returns
    -------
    ggplot
        A plotnine ggplot object that can be further customized.
    """
    if isinstance(result, DDDAggResult):
        if result.aggregation_type not in ("group", "calendar"):
            raise ValueError(
                f"plot_agg requires group or calendar aggregation, got {result.aggregation_type}. "
                f"Use plot_event_study for eventstudy aggregation."
            )
        df = dddaggresult_to_polars(result)
        is_ddd = True
    elif isinstance(result, AGGTEResult):
        if result.aggregation_type not in ("group", "calendar"):
            raise ValueError(
                f"plot_agg requires group or calendar aggregation, got {result.aggregation_type}. "
                f"Use plot_event_study for dynamic aggregation."
            )
        df = aggteresult_to_polars(result)
        is_ddd = False
    else:
        raise TypeError(f"plot_agg requires AGGTEResult or DDDAggResult, got {type(result).__name__}")

    if result.aggregation_type == "group":
        default_xlab = "Treatment Cohort"
        default_title = "DDD Effects by Treatment Cohort" if is_ddd else "Effects by Treatment Cohort"
    else:
        default_xlab = "Calendar Time"
        default_title = "DDD Effects by Calendar Time" if is_ddd else "Effects by Calendar Time"

    plot = ggplot(df, aes(x="event_time", y="att"))

    if show_ci:
        plot = plot + geom_errorbar(
            aes(ymin="ci_lower", ymax="ci_upper"),
            width=0.2,
            size=0.8,
            color=COLORS["post_treatment"],
        )

    if ref_line is not None:
        plot = plot + geom_hline(yintercept=ref_line, linetype="dashed", color="#7f8c8d", alpha=0.7)

    x_breaks = sorted(df["event_time"].unique().to_list())

    plot = (
        plot
        + geom_line(color=COLORS["line"], size=0.8, alpha=0.6, linetype="dotted")
        + geom_point(color=COLORS["post_treatment"], size=3.5)
        + scale_x_continuous(breaks=x_breaks)
        + labs(
            x=xlab or default_xlab,
            y=ylab or "ATT",
            title=title or default_title,
        )
        + theme_minimal()
    )

    return plot


def plot_dose_response(
    result: DoseResult,
    effect_type: Literal["att", "acrt"] = "att",
    show_ci: bool = True,
    ref_line: float | None = 0,
    xlab: str | None = None,
    ylab: str | None = None,
    title: str | None = None,
    **_kwargs: Any,
) -> ggplot:
    """Plot dose-response function for continuous treatment.

    Parameters
    ----------
    result : DoseResult
        Continuous treatment dose-response result.
    effect_type : {'att', 'acrt'}, default='att'
        Type of effect to plot:
        - 'att': Average Treatment Effect on Treated
        - 'acrt': Average Causal Response on Treated (marginal effect)
    show_ci : bool, default=True
        Whether to show confidence bands.
    ref_line : float or None, default=0
        Y-value for reference line. Set to None to hide.
    xlab : str, optional
        X-axis label. Defaults to "Dose".
    ylab : str, optional
        Y-axis label. Defaults based on effect_type.
    title : str, optional
        Plot title. Defaults based on effect_type.

    Returns
    -------
    ggplot
        A plotnine ggplot object that can be further customized.
    """
    df = doseresult_to_polars(result, effect_type=effect_type)

    default_ylabel = "ATT(d)" if effect_type == "att" else "ACRT(d)"
    default_title = f"Dose-Response: {default_ylabel}"

    line_color = "#2c3e50"
    fill_color = "#95a5a6"

    plot = ggplot(df, aes(x="dose", y="effect"))

    if show_ci:
        plot = plot + geom_ribbon(
            aes(ymin="ci_lower", ymax="ci_upper"),
            fill=fill_color,
            alpha=0.25,
        )

    if ref_line is not None:
        plot = plot + geom_hline(yintercept=ref_line, linetype="dashed", color="#7f8c8d", alpha=0.7)

    plot = (
        plot
        + geom_line(color=line_color, size=1.2)
        + geom_point(color=line_color, size=2.5)
        + labs(
            x=xlab or "Dose",
            y=ylab or default_ylabel,
            title=title or default_title,
        )
        + theme_minimal()
    )

    return plot


def plot_sensitivity(
    result: HonestDiDResult,
    ref_line: float | None = 0,
    xlab: str | None = None,
    ylab: str | None = None,
    title: str | None = None,
    **_kwargs: Any,
) -> ggplot:
    """Create sensitivity analysis plot for HonestDiD results.

    Parameters
    ----------
    result : HonestDiDResult
        Honest DiD sensitivity analysis result.
    ref_line : float or None, default=0
        Y-value for reference line. Set to None to hide.
    xlab : str, optional
        X-axis label. Defaults based on sensitivity type.
    ylab : str, optional
        Y-axis label. Defaults to "Confidence Interval".
    title : str, optional
        Plot title. Defaults based on sensitivity type.

    Returns
    -------
    ggplot
        A plotnine ggplot object that can be further customized.
    """
    df = honestdid_to_polars(result)

    is_smoothness = result.sensitivity_type == "smoothness"
    default_xlab = "M" if is_smoothness else r"$\bar{M}$"
    default_title = f"Sensitivity Analysis ({result.sensitivity_type.replace('_', ' ').title()})"

    method_colors = {
        "Original": COLORS["original"],
        "FLCI": COLORS["flci"],
        "Conditional": COLORS["conditional"],
        "C-F": COLORS["c_f"],
        "C-LF": COLORS["c_lf"],
    }

    methods = df["method"].unique().to_list()
    available_colors = {m: method_colors.get(m, "#34495e") for m in methods}

    n_methods = len(methods)
    dodge_width = 0.05 * (df["param_value"].max() - df["param_value"].min()) if n_methods > 1 else 0

    plot = (
        ggplot(df, aes(x="param_value", y="midpoint", color="method"))
        + geom_point(size=3, position=position_dodge(width=dodge_width))
        + geom_errorbar(
            aes(ymin="lb", ymax="ub"),
            width=0.02 * (df["param_value"].max() - df["param_value"].min()),
            size=0.8,
            position=position_dodge(width=dodge_width),
        )
        + scale_color_manual(values=available_colors, name="Method")
        + labs(
            x=xlab or default_xlab,
            y=ylab or "Confidence Interval",
            title=title or default_title,
        )
        + theme_minimal()
        + theme(legend_position="bottom")
    )

    if ref_line is not None:
        plot = plot + geom_hline(yintercept=ref_line, linetype="solid", color="black", alpha=0.4)

    return plot


def plot_multiplegt(
    result: DIDInterResult,
    show_ci: bool = True,
    ref_line: float | None = 0,
    xlab: str | None = None,
    ylab: str | None = None,
    title: str | None = None,
    **_kwargs: Any,
) -> ggplot:
    """Create event study plot for intertemporal treatment effects.

    Parameters
    ----------
    result : DIDInterResult
        Intertemporal treatment effects result from did_multiplegt().
    show_ci : bool, default=True
        Whether to show confidence intervals as error bars.
    ref_line : float or None, default=0
        Y-value for reference line. Set to None to hide.
    xlab : str, optional
        X-axis label. Defaults to "Horizon".
    ylab : str, optional
        Y-axis label. Defaults to "Effect".
    title : str, optional
        Plot title. Defaults to "Intertemporal Treatment Effects".

    Returns
    -------
    ggplot
        A plotnine ggplot object that can be further customized.
    """
    df = didinterresult_to_polars(result)
    x_breaks = sorted(df["horizon"].unique().to_list())

    plot = ggplot(df, aes(x="horizon", y="att"))

    if show_ci:
        plot = plot + geom_errorbar(
            aes(ymin="ci_lower", ymax="ci_upper", color="treatment_status"),
            width=0.2,
            size=0.8,
        )

    if ref_line is not None:
        plot = plot + geom_hline(yintercept=ref_line, linetype="dashed", color="#7f8c8d", alpha=0.7)

    plot = plot + geom_vline(xintercept=0, linetype="dotted", color="#7f8c8d", alpha=0.7)

    plot = (
        plot
        + geom_line(color=COLORS["line"], size=0.8, alpha=0.6, linetype="dotted")
        + geom_point(aes(color="treatment_status"), size=3.5)
        + scale_color_manual(
            values={"Pre": COLORS["pre_treatment"], "Post": COLORS["post_treatment"]},
            limits=["Pre", "Post"],
            name="Treatment Status",
        )
        + scale_x_continuous(breaks=x_breaks)
        + labs(
            x=xlab or "Horizon",
            y=ylab or "Effect",
            title=title or "Intertemporal Treatment Effects",
        )
        + theme_minimal()
        + theme(legend_position="bottom")
    )

    return plot
