"""Frozen metrics for the independent actor--critic oral-B-v2 branch."""
import math

import torch

from sdil.bci_metrics import (
    grouped_classification_accuracy,
    grouped_regression_correlation,
    pearson,
)


def annotate_events(events, day, final_success, episode_offset, horizon):
    """Attach stable groups, outcomes, and assay horizon to event records."""
    final_success = torch.as_tensor(final_success, dtype=torch.bool)
    episode_ids = (
        torch.arange(final_success.numel(), dtype=torch.long)
        + episode_offset
    )
    for event in events:
        event["day"] = day
        event["episode_ids"] = episode_ids.clone()
        event["final_success"] = final_success.clone()
        event["horizon"] = horizon
    return events


def flatten_active(events):
    names = (
        "soma",
        "previous_soma",
        "raw_apical",
        "innovation",
        "abs_error",
        "previous_abs_error",
        "td_innovation",
        "value_prediction",
        "terminal",
        "outcome",
    )
    values = {name: [] for name in names}
    values.update({"groups": [], "day": [], "horizon": []})
    for event in events:
        active = event["active"].bool()
        if not bool(active.any()):
            continue
        count = int(active.sum())
        for name in names:
            values[name].append(event[name][active].double())
        values["groups"].append(event["episode_ids"][active].long())
        values["day"].append(
            torch.full((count,), int(event["day"]), dtype=torch.long)
        )
        values["horizon"].append(
            torch.full((count,), int(event["horizon"]), dtype=torch.long)
        )
    if not values["soma"]:
        raise ValueError("no active oral-B-v2 events")
    return {key: torch.cat(chunks, dim=0) for key, chunks in values.items()}


def population_correlations(flat, left, right, mask=None):
    if mask is None:
        mask = torch.ones(flat[left].shape[0], dtype=torch.bool)
    return [
        pearson(flat[left][mask, cell], flat[right][mask, cell])
        for cell in range(flat[left].shape[1])
    ]


def terminal_population(events):
    """Return exactly one terminal record per annotated episode."""
    records = {}
    for event in events:
        for row, group in enumerate(event["episode_ids"].tolist()):
            if not bool(event["terminal"][row]):
                continue
            if group in records:
                raise ValueError("multiple terminal events for one episode")
            records[group] = {
                "innovation": event["innovation"][row].double(),
                "previous_soma": event["previous_soma"][row].double(),
                "soma": event["soma"][row].double(),
                "td_innovation": event["td_innovation"][row].double(),
                "value_prediction": event["value_prediction"][row].double(),
                "outcome": bool(event["outcome"][row]),
                "horizon": int(event["horizon"]),
            }
    expected = {
        int(group)
        for event in events
        for group in event["episode_ids"].tolist()
    }
    if set(records) != expected:
        raise ValueError("every annotated episode needs one terminal event")
    groups = sorted(records)
    return {
        "innovation": torch.stack([
            records[group]["innovation"] for group in groups
        ]),
        "previous_soma": torch.stack([
            records[group]["previous_soma"] for group in groups
        ]),
        "soma": torch.stack([
            records[group]["soma"] for group in groups
        ]),
        "td_innovation": torch.stack([
            records[group]["td_innovation"] for group in groups
        ]),
        "value_prediction": torch.stack([
            records[group]["value_prediction"] for group in groups
        ]),
        "outcome": torch.tensor([
            records[group]["outcome"] for group in groups
        ], dtype=torch.bool),
        "horizon": torch.tensor([
            records[group]["horizon"] for group in groups
        ], dtype=torch.long),
        "groups": torch.tensor(groups, dtype=torch.long),
    }


def _paired_terminal_populations(mode_events):
    populations = {
        name: terminal_population(events)
        for name, events in mode_events.items()
    }
    reference = populations["intact"]
    for name, population in populations.items():
        for key in ("outcome", "horizon", "groups"):
            if not torch.equal(reference[key], population[key]):
                raise ValueError(
                    f"challenge mode {name} is not paired on {key}"
                )
    return populations


def _mean_by_label(values, labels, label):
    selected = values[labels == label]
    if not selected.numel():
        return 0.0
    return selected.double().mean().item()


def signature_metrics_v2(train_events, challenge_mode_events, cfg, role):
    """Measure nonterminal innovation and terminal outcome surprise."""
    train = flatten_active(train_events)
    role = torch.as_tensor(role, dtype=torch.float64)
    nonterminal = ~train["terminal"].bool()
    if not bool(nonterminal.any()):
        raise ValueError("v2 training needs nonterminal events")

    soma_corr = population_correlations(
        train, "innovation", "soma", nonterminal
    )
    raw_corr = population_correlations(
        train, "raw_apical", "soma", nonterminal
    )

    event_decoder = []
    confidence_corr = []
    for cell in range(cfg.n_plus + cfg.n_minus):
        keep = [
            index for index in range(cfg.n_neurons)
            if index != cell
        ]
        labels = train["innovation"][nonterminal, cell] > 0
        accuracy, distance = grouped_classification_accuracy(
            train["previous_soma"][nonterminal][:, keep],
            labels,
            train["groups"][nonterminal],
        )
        event_decoder.append(accuracy)
        confidence_corr.append(
            pearson(
                distance,
                train["innovation"][nonterminal, cell],
            )
        )

    difference = (
        train["innovation"][nonterminal, :cfg.n_plus].mean(1)
        - train["innovation"][
            nonterminal,
            cfg.n_plus:cfg.n_plus + cfg.n_minus,
        ].mean(1)
    )
    improvement = (
        train["previous_abs_error"][nonterminal]
        - train["abs_error"][nonterminal]
    )
    improving = improvement > 0
    worsening = improvement < 0
    sign_index = 0.0
    if bool(improving.any()) and bool(worsening.any()):
        sign_index = 0.5 * (
            difference[improving].mean()
            - difference[worsening].mean()
        ).item()
    aligned = train["innovation"][nonterminal] @ role
    velocity_corr = grouped_regression_correlation(
        aligned, improvement, train["groups"][nonterminal]
    )
    error_corr = grouped_regression_correlation(
        aligned,
        train["abs_error"][nonterminal],
        train["groups"][nonterminal],
    )

    populations = _paired_terminal_populations(challenge_mode_events)
    intact = populations["intact"]
    critic_lesion = populations["acute_critic_lesion"]
    outcome_lesion = populations["acute_outcome_lesion"]
    labels = intact["outcome"]
    groups = intact["groups"]
    residual_accuracy, _ = grouped_classification_accuracy(
        intact["innovation"], labels, groups
    )
    previous_soma_accuracy, _ = grouped_classification_accuracy(
        intact["previous_soma"], labels, groups
    )
    outcome_lesion_accuracy, _ = grouped_classification_accuracy(
        outcome_lesion["innovation"], labels, groups
    )

    scores = {
        name: population["innovation"] @ role
        for name, population in populations.items()
    }
    intact_separation = (
        _mean_by_label(scores["intact"], labels, True)
        - _mean_by_label(scores["intact"], labels, False)
    )
    outcome_lesion_separation = (
        _mean_by_label(scores["acute_outcome_lesion"], labels, True)
        - _mean_by_label(
            scores["acute_outcome_lesion"], labels, False
        )
    )
    critic_contribution = (
        scores["acute_critic_lesion"] - scores["intact"]
    )
    critic_prediction_corr = pearson(
        critic_contribution, intact["value_prediction"]
    )

    success_fraction = labels.double().mean().item()
    horizon_success = {}
    for horizon in sorted(set(intact["horizon"].tolist())):
        selected = intact["horizon"] == horizon
        horizon_success[str(horizon)] = (
            labels[selected].double().mean().item()
        )

    values = {
        "mean_abs_residual_soma_corr":
            sum(abs(value) for value in soma_corr) / len(soma_corr),
        "mean_abs_raw_soma_corr":
            sum(abs(value) for value in raw_corr) / len(raw_corr),
        "raw_minus_residual_abs_soma_corr": (
            sum(abs(value) for value in raw_corr) / len(raw_corr)
            - sum(abs(value) for value in soma_corr) / len(soma_corr)
        ),
        "surrounding_event_decoder_balanced_acc":
            sum(event_decoder) / len(event_decoder),
        "decoder_distance_residual_corr":
            sum(confidence_corr) / len(confidence_corr),
        "causal_role_sign_inversion_index": sign_index,
        "role_aligned_velocity_cv_corr": velocity_corr,
        "role_aligned_error_cv_corr": error_corr,
        "velocity_minus_error_abs_cv_corr":
            abs(velocity_corr) - abs(error_corr),
        "terminal_residual_outcome_balanced_acc": residual_accuracy,
        "terminal_previous_soma_outcome_balanced_acc":
            previous_soma_accuracy,
        "terminal_residual_minus_previous_soma_acc":
            residual_accuracy - previous_soma_accuracy,
        "acute_outcome_lesion_outcome_balanced_acc":
            outcome_lesion_accuracy,
        "terminal_role_aligned_outcome_separation": intact_separation,
        "acute_outcome_lesion_role_aligned_separation":
            outcome_lesion_separation,
        "terminal_outcome_separation_drop_under_acute_lesion":
            intact_separation - outcome_lesion_separation,
        "mean_critic_expectedness_contribution":
            critic_contribution.mean().item(),
        "critic_contribution_value_prediction_corr":
            critic_prediction_corr,
        "challenge_success_fraction": success_fraction,
        "challenge_success_count": int(labels.sum()),
        "challenge_failure_count": int((~labels).sum()),
        "challenge_episodes": int(labels.numel()),
        "nonterminal_training_events": int(nonterminal.sum()),
        "terminal_training_events": int(train["terminal"].sum()),
        "horizon_success_fraction": horizon_success,
    }
    flat_values = {
        key: value for key, value in values.items()
        if isinstance(value, (int, float))
    }
    if not all(math.isfinite(value) for value in flat_values.values()):
        raise ValueError("non-finite oral-B-v2 signature metric")
    return values