"""Predetermined metrics for the preregistered continuous-BCI signatures."""
import math

import torch


def pearson(x, y):
    x = torch.as_tensor(x, dtype=torch.float64).flatten()
    y = torch.as_tensor(y, dtype=torch.float64).flatten()
    finite = torch.isfinite(x) & torch.isfinite(y)
    x, y = x[finite], y[finite]
    if x.numel() < 3:
        return 0.0
    x = x - x.mean()
    y = y - y.mean()
    denominator = x.square().sum().sqrt() * y.square().sum().sqrt()
    if denominator <= 1e-15:
        return 0.0
    return (x @ y / denominator).item()


def balanced_accuracy(labels, predictions):
    labels = torch.as_tensor(labels, dtype=torch.bool)
    predictions = torch.as_tensor(predictions, dtype=torch.bool)
    scores = []
    for value in (False, True):
        mask = labels == value
        if bool(mask.any()):
            scores.append((predictions[mask] == labels[mask]).double().mean())
    return torch.stack(scores).mean().item() if scores else 0.5


def grouped_linear_predictions(features, target, groups, *, classification,
                               n_folds=5, ridge=1e-3):
    """Out-of-group-fold linear scores with train-fold standardization."""
    x = torch.as_tensor(features, dtype=torch.float64)
    y = torch.as_tensor(target, dtype=torch.float64).flatten()
    groups = torch.as_tensor(groups, dtype=torch.long).flatten()
    if x.ndim == 1:
        x = x.unsqueeze(1)
    if x.shape[0] != y.numel() or groups.numel() != y.numel():
        raise ValueError("features, targets, and groups must have matching rows")
    if x.shape[0] < n_folds * 2:
        return torch.zeros_like(y)
    prediction = torch.zeros_like(y)
    for fold in range(n_folds):
        test = groups.remainder(n_folds) == fold
        train = ~test
        if not bool(test.any()) or train.sum() <= x.shape[1] + 1:
            continue
        train_x, train_y = x[train], y[train]
        mean = train_x.mean(0)
        scale = train_x.std(0, unbiased=False).clamp_min(1e-8)
        train_x = (train_x - mean) / scale
        test_x = (x[test] - mean) / scale
        train_x = torch.cat((train_x, torch.ones(train_x.shape[0], 1)), dim=1)
        test_x = torch.cat((test_x, torch.ones(test_x.shape[0], 1)), dim=1)
        if classification:
            binary = train_y > 0.5
            positives = binary.sum().item()
            negatives = binary.numel() - positives
            if positives == 0 or negatives == 0:
                prediction[test] = 1.0 if positives else -1.0
                continue
            signed_y = binary.double().mul(2).sub(1)
            weights = torch.where(
                binary,
                torch.full_like(train_y, binary.numel() / (2.0 * positives)),
                torch.full_like(train_y, binary.numel() / (2.0 * negatives)))
            root_weight = weights.sqrt()
            fit_x = train_x * root_weight.unsqueeze(1)
            fit_y = signed_y * root_weight
        else:
            fit_x, fit_y = train_x, train_y
        penalty = torch.eye(fit_x.shape[1], dtype=torch.float64) * ridge
        penalty[-1, -1] = 0.0
        gram = fit_x.t() @ fit_x + penalty
        rhs = fit_x.t() @ fit_y
        try:
            coefficient = torch.linalg.solve(gram, rhs)
        except RuntimeError:
            coefficient = torch.linalg.pinv(gram) @ rhs
        prediction[test] = test_x @ coefficient
    return prediction


def grouped_classification_accuracy(features, labels, groups):
    labels = torch.as_tensor(labels, dtype=torch.bool)
    if labels.numel() < 10 or labels.all() or (~labels).all():
        return 0.5, torch.zeros(labels.numel(), dtype=torch.float64)
    score = grouped_linear_predictions(
        features, labels.double(), groups, classification=True)
    return balanced_accuracy(labels, score > 0), score


def grouped_regression_correlation(features, target, groups):
    prediction = grouped_linear_predictions(
        features, target, groups, classification=False)
    return pearson(prediction, target)


def annotate_day_events(events, day, final_success, episode_offset):
    """Attach stable episode groups and final outcomes to collected 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()
    return events


def flatten_active(events):
    fields = {
        "soma": [], "previous_soma": [], "raw": [], "innovation": [],
        "abs_error": [], "previous_abs_error": [], "groups": [], "day": [],
        "outcome": [],
    }
    for event in events:
        active = event["active"].bool()
        if not bool(active.any()):
            continue
        count = active.sum().item()
        fields["soma"].append(event["base_soma"][active].double())
        fields["previous_soma"].append(event["previous_soma"][active].double())
        fields["raw"].append(event["raw_apical"][active].double())
        fields["innovation"].append(event["innovation"][active].double())
        fields["abs_error"].append(event["abs_error"][active].double())
        fields["previous_abs_error"].append(
            event["previous_abs_error"][active].double())
        fields["groups"].append(event["episode_ids"][active].long())
        fields["day"].append(torch.full((count,), int(event["day"]), dtype=torch.long))
        fields["outcome"].append(event["final_success"][active].bool())
    if not fields["soma"]:
        raise ValueError("no active BCI events were collected")
    return {key: torch.cat(value, dim=0) for key, value in fields.items()}


def episode_population_features(events, field):
    """Mean of the last four active events for every episode."""
    values = {}
    outcomes = {}
    for event in events:
        for row, group in enumerate(event["episode_ids"].tolist()):
            if not bool(event["active"][row]):
                continue
            values.setdefault(group, []).append(event[field][row].double())
            outcomes[group] = bool(event["final_success"][row])
    groups = sorted(values)
    features = torch.stack([
        torch.stack(values[group][-4:]).mean(0) for group in groups])
    labels = torch.tensor([outcomes[group] for group in groups], dtype=torch.bool)
    return features, labels, torch.tensor(groups, dtype=torch.long)


def population_correlations(flat, left, right):
    return [pearson(flat[left][:, cell], flat[right][:, cell])
            for cell in range(flat[left].shape[1])]


def signature_metrics(train_events, evaluation_events, cfg, role):
    train = flatten_active(train_events)
    role = torch.as_tensor(role, dtype=torch.float64)
    soma_corr = population_correlations(train, "innovation", "soma")
    raw_corr = population_correlations(train, "raw", "soma")

    # Surrounding-network prediction of each causal cell's amplification sign.
    event_decoder, confidence_corr = [], []
    causal_cells = range(cfg.n_plus + cfg.n_minus)
    for cell in causal_cells:
        keep = [index for index in range(cfg.n_neurons) if index != cell]
        labels = train["innovation"][:, cell] > 0
        score, distance = grouped_classification_accuracy(
            train["previous_soma"][:, keep], labels, train["groups"])
        event_decoder.append(score)
        confidence_corr.append(pearson(distance, train["innovation"][:, cell]))

    diff = (train["innovation"][:, :cfg.n_plus].mean(1)
            - train["innovation"][:, cfg.n_plus:cfg.n_plus + cfg.n_minus].mean(1))
    improvement = train["previous_abs_error"] - train["abs_error"]
    improving = improvement > 0
    worsening = improvement < 0
    sign_index = 0.0
    if bool(improving.any()) and bool(worsening.any()):
        sign_index = 0.5 * (diff[improving].mean() - diff[worsening].mean()).item()

    role_aligned = train["innovation"] @ role
    velocity_corr = grouped_regression_correlation(
        role_aligned, improvement, train["groups"])
    error_corr = grouped_regression_correlation(
        role_aligned, train["abs_error"], train["groups"])

    residual_features, outcomes, outcome_groups = episode_population_features(
        evaluation_events, "innovation")
    soma_features, soma_outcomes, soma_groups = episode_population_features(
        evaluation_events, "base_soma")
    if not torch.equal(outcomes, soma_outcomes) or not torch.equal(outcome_groups, soma_groups):
        raise ValueError("residual and soma outcome populations are not paired")
    residual_outcome, _ = grouped_classification_accuracy(
        residual_features, outcomes, outcome_groups)
    soma_outcome, _ = grouped_classification_accuracy(
        soma_features, outcomes, outcome_groups)

    early = train["day"] <= 2
    late = train["day"] >= cfg.days - 3
    if not bool(early.any()) or not bool(late.any()):
        longitudinal = 0.0
    else:
        early_residual = train["innovation"][early].mean(0)
        activity_change = (train["soma"][late].mean(0)
                           - train["soma"][early].mean(0))
        longitudinal = pearson(early_residual, activity_change)

    values = {
        "mean_abs_residual_soma_corr": sum(abs(v) for v in soma_corr) / len(soma_corr),
        "mean_abs_raw_soma_corr": sum(abs(v) for v in raw_corr) / len(raw_corr),
        "raw_minus_residual_abs_soma_corr": (
            sum(abs(v) for v in raw_corr) / len(raw_corr)
            - sum(abs(v) for v 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),
        "residual_outcome_decoder_balanced_acc": residual_outcome,
        "soma_outcome_decoder_balanced_acc": soma_outcome,
        "residual_minus_soma_outcome_acc": residual_outcome - soma_outcome,
        "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),
        "early_residual_late_activity_change_corr": longitudinal,
        "active_training_events": int(train["soma"].shape[0]),
        "evaluation_episodes": int(outcomes.numel()),
        "evaluation_success_fraction": outcomes.double().mean().item(),
    }
    if not all(math.isfinite(value) for value in values.values()):
        raise ValueError("non-finite BCI signature metric")
    return values