from __future__ import annotations

import pickle

import numpy as np
import pandas as pd
import torch

from biolmnet.artifacts import load_bundle, save_bundle
from biolmnet.data import (
    BranchPriors,
    PreparedWorkspace,
    attach_embeddings_and_pathways,
    build_biological_mask,
    deterministic_gene_embeddings,
)
from biolmnet.model import BioMaskedLinear, BioLMNet
from biolmnet.training import Hyperparameters, predict, train


def test_biological_mask_uses_pdi_and_undirected_ppi() -> None:
    pdi = pd.DataFrame(
        {"TF": ["A", "C", "outside"], "Target": ["B", "D", "A"]}
    )
    ppi = pd.DataFrame(
        {
            "protein1": ["A", "X", "B"],
            "protein2": ["X", "C", "Y"],
            "combined_score": [950, 950, 710],
        }
    )
    branch = build_biological_mask(["A", "B", "C", "D"], pdi, ppi)

    assert {"B", "D", "X"}.issubset(branch.hidden_genes)
    x_index = branch.hidden_genes.index("X")
    assert branch.biological_mask[0, x_index] > 0
    assert branch.biological_mask[2, x_index] > 0
    assert branch.pdi_edges == 2


def test_masked_linear_disconnects_unlisted_weights() -> None:
    mask = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
    layer = BioMaskedLinear(mask, bias=False)
    values = torch.tensor([[2.0, 3.0]])
    baseline = layer(values).detach().clone()
    with torch.no_grad():
        layer.weight[0, 1] = 10_000
        layer.weight[1, 0] = -10_000
    changed = layer(values).detach()
    assert torch.allclose(baseline, changed)


def _tiny_workspace() -> PreparedWorkspace:
    rng = np.random.default_rng(7)
    sample_count = 48
    labels = np.repeat(np.array([0, 1]), sample_count // 2)
    gene_values = rng.normal(size=(sample_count, 4)).astype(np.float32)
    dna_values = rng.normal(size=(sample_count, 4)).astype(np.float32)
    gene_values[:, 0] += labels * 1.5
    dna_values[:, 1] -= labels * 1.2
    input_genes = ["A", "B", "C", "D"]
    hidden_genes = ["A", "B", "C"]
    biological = np.array(
        [
            [1.0, 1.0, 0.0],
            [0.0, 1.0, 1.0],
            [1.0, 0.0, 1.0],
            [0.0, 1.0, 0.0],
        ],
        dtype=np.float32,
    )
    pathway_mapping = pd.DataFrame(
        {
            "SYMBOL": ["A", "B", "B", "C"],
            "PathwayID": ["hsa1", "hsa1", "hsa2", "hsa2"],
        }
    )
    embeddings = deterministic_gene_embeddings(hidden_genes, dimensions=8)

    def branch() -> BranchPriors:
        value = BranchPriors(
            input_genes=input_genes.copy(),
            hidden_genes=hidden_genes.copy(),
            biological_mask=biological.copy(),
            pdi_edges=3,
            ppi_edges=4,
        )
        attach_embeddings_and_pathways(
            value,
            embeddings,
            pathway_mapping,
            precomputed_significant=True,
        )
        return value

    return PreparedWorkspace(
        gene_expression=gene_values,
        dna_methylation=dna_values,
        labels=labels,
        label_names=["control", "case"],
        gene_branch=branch(),
        dna_branch=branch(),
        source_name="unit test",
    )


def test_model_forward_probabilistic_shape() -> None:
    workspace = _tiny_workspace()
    gene = workspace.gene_branch
    dna = workspace.dna_branch
    model = BioLMNet(
        torch.from_numpy(gene.biological_mask),
        torch.from_numpy(dna.biological_mask),
        torch.from_numpy(gene.embeddings),
        torch.from_numpy(dna.embeddings),
        torch.from_numpy(gene.pathway_mask),
        torch.from_numpy(dna.pathway_mask),
        n_classes=2,
        projection_dim=4,
        fusion_dim=3,
        dropout=0.0,
    )
    logits = model(torch.randn(5, 4), torch.randn(5, 4))
    assert logits.shape == (5, 2)
    assert torch.allclose(
        model.gene_branch.pathway_attention.attention_weights().sum(dim=0),
        torch.ones(2),
    )


def test_training_artifact_roundtrip(tmp_path) -> None:
    workspace = _tiny_workspace()
    result = train(
        workspace,
        Hyperparameters(
            epochs=3,
            batch_size=8,
            projection_dim=4,
            fusion_dim=3,
            dropout=0.0,
            early_stopping_patience=3,
        ),
    )
    restored_from_process_boundary = pickle.loads(pickle.dumps(result.bundle))
    assert restored_from_process_boundary.label_names == ["control", "case"]
    artifact = save_bundle(result.bundle, tmp_path / "model.zip")
    restored = load_bundle(artifact)
    gene_frame = pd.DataFrame(
        workspace.gene_expression[:5], columns=restored.gene_features
    )
    dna_frame = pd.DataFrame(
        workspace.dna_methylation[:5], columns=restored.dna_features
    )
    before = predict(gene_frame, dna_frame, result.bundle)
    after = predict(gene_frame, dna_frame, restored)
    probability_columns = [column for column in before if column.startswith("P(")]
    np.testing.assert_allclose(
        before[probability_columns].to_numpy(),
        after[probability_columns].to_numpy(),
        atol=1e-6,
    )


def test_zerogpu_duration_estimator_is_bounded_and_scales() -> None:
    from app import estimate_training_duration

    workspace = _tiny_workspace()
    common = (workspace, 16, 0.001, 0.01, 0.3, 64, 12, 0.2, "Adam", True)
    short = estimate_training_duration(common[0], 10, *common[1:])
    long = estimate_training_duration(common[0], 200, *common[1:])
    assert 30 <= short <= 300
    assert short <= long <= 300