from __future__ import annotations

import pandas as pd
import pytest

from phospy.matrices import build_site_matrix
from phospy.validation.errors import TableSchemaError


def test_build_site_matrix_creates_site_ids_and_deduplicates_by_mean() -> None:
    df = pd.DataFrame(
        {
            "gene_p_site": ["PRKACA_S339", "PRKACA_S339", "BTK_Y551"],
            "centralized_sequence": ["AAAAAA", "BBBBBB", "CCCCCC"],
            "phospho_corrected_1": [1.0, 10.0, 3.0],
            "phospho_corrected_2": [1.0, 10.0, 3.0],
            "phospho_corrected_3": [1.0, 10.0, 3.0],
            "phospho_corrected_4": [1.0, 10.0, 3.0],
            "phospho_corrected_5": [1.0, 10.0, 3.0],
            "phospho_corrected_6": [1.0, 10.0, 3.0],
        }
    )
    phosr_input, matrix, sequences = build_site_matrix(
        df=df,
        gene_p_site_col="gene_p_site",
        sequence_col="centralized_sequence",
        value_cols=[
            "phospho_corrected_1",
            "phospho_corrected_2",
            "phospho_corrected_3",
            "phospho_corrected_4",
            "phospho_corrected_5",
            "phospho_corrected_6",
        ],
    )
    assert "PRKACA;S339;" in matrix.index
    assert matrix.loc["PRKACA;S339;", "phospho_corrected_1"] == 10.0
    assert sequences.loc["PRKACA;S339;"] == "BBBBBB"
    assert phosr_input.shape[0] == 2


def test_build_site_matrix_exposes_row_drop_stats() -> None:
    df = pd.DataFrame(
        {
            "gene_p_site": ["PRKACA_S339", "BTK_Y551", "BTK_Y551"],
            "centralized_sequence": ["AAAAAA", None, "CCCCCC"],
            "phospho_corrected_1": [1.0, 2.0, 3.0],
            "phospho_corrected_2": [1.0, 2.0, None],
        }
    )
    phosr_input, matrix, sequences = build_site_matrix(
        df=df,
        gene_p_site_col="gene_p_site",
        sequence_col="centralized_sequence",
        value_cols=["phospho_corrected_1", "phospho_corrected_2"],
    )

    stats = phosr_input.attrs["row_drop_stats"]
    assert stats["input_rows"] == 3
    assert stats["dropped_missing_sequence"] == 1
    assert stats["dropped_incomplete_values"] == 1
    assert stats["retained_rows"] == 1
    assert matrix.attrs["row_drop_stats"] == stats
    assert sequences.attrs["row_drop_stats"] == stats


def test_build_site_matrix_raises_table_schema_error_for_malformed_gene_p_site() -> (
    None
):
    df = pd.DataFrame(
        {
            "gene_p_site": ["PRKACA", "BTK_Y551"],
            "centralized_sequence": ["AAAAAA", "CCCCCC"],
            "phospho_corrected_1": [1.0, 3.0],
            "phospho_corrected_2": [1.0, 3.0],
        }
    )

    with pytest.raises(
        TableSchemaError, match="must contain values in the form GENE_SITE"
    ):
        build_site_matrix(
            df=df,
            gene_p_site_col="gene_p_site",
            sequence_col="centralized_sequence",
            value_cols=["phospho_corrected_1", "phospho_corrected_2"],
        )


def test_site_matrix_builder_reports_row_drop_diagnostics_when_all_rows_are_dropped() -> (
    None
):
    from phospy.site_matrix_builder import SiteMatrixBuilder

    df = pd.DataFrame(
        {
            "gene_p_site": ["PRKACA_S339", "BTK_Y551", "SRC_Y419"],
            "centralized_sequence": [None, "CCCCCC", "DDDDDD"],
            "phospho_corrected_1": [1.0, None, 3.0],
            "phospho_corrected_2": [1.0, 2.0, None],
        }
    )

    builder = SiteMatrixBuilder(
        value_cols=["phospho_corrected_1", "phospho_corrected_2"]
    )

    with pytest.raises(TableSchemaError, match="row-drop diagnostics") as exc_info:
        builder.build(df)

    message = str(exc_info.value)
    assert "site matrix must contain at least one phosphosite row" in message
    assert "dropped_missing_sequence=1" in message
    assert "dropped_incomplete_values=2" in message
    assert "deduplicated_site_rows=0" in message
    assert "other_dropped_rows=0" in message
    assert "retained_rows=0" in message
