import unittest

import numpy as np
import pandas as pd

import biological_contract as b


def fixture():
    types = ["KC-a", "KC-b", "PAM15", "PPL103", "MBON01", "MBON12", "other-DAN"]
    classes = ["Kenyon_Cell", "Kenyon_Cell", "DAN", "DAN", "MBON", "MBON", "DAN"]
    instances = ["KC-a_L", "KC-b_R", "PAM15(y5B'2a)_L", "PPL103(y2a'1)_R",
                 "MBON01(y5B'2a)_L", "MBON12(y2a'1)_R", "other-DAN"]
    meta = pd.DataFrame(dict(neuron_idx=np.arange(7, dtype=np.int32),
        body_id=np.arange(7, dtype=np.uint64) + np.uint64(2**63),
        annotation_type=pd.Series(types, dtype=object), annotation_class=pd.Series(classes, dtype=object),
        annotation_instance=pd.Series(instances, dtype=object)))
    dan = meta.iloc[[2, 3, 6]].drop(columns="annotation_class").copy()
    mbon = meta.iloc[[4, 5]].drop(columns="annotation_class").copy()
    dan["mb_compartment_verbatim"] = ["y5B'2a", "y2a'1", None]
    dan["mb_compartment_mapping_status"] = ["explicit_single_parenthesized_token"] * 2 + ["no_parenthesized_token"]
    mbon["mb_compartment_verbatim"] = ["y5B'2a", "y2a'1"]
    mbon["mb_compartment_mapping_status"] = "explicit_single_parenthesized_token"
    src = np.array([0, 0, 1, 1, 2, 3], np.int32)
    dst = np.array([4, 5, 4, 5, 4, 5], np.int32)
    raw = np.array([4, 8, 12, 16, 20, 24], np.int64)
    def table(ids):
        rows = []
        for i in ids:
            row = dict(edge_id=i, raw_synapse_count=int(raw[i]))
            for label, node in (("source", src[i]), ("target", dst[i])):
                for key in ("neuron_idx", "body_id", "annotation_type", "annotation_instance"):
                    row[label + "_" + key] = meta.iloc[node][key]
            row["target_compartment_verbatim"] = "y5B'2a" if dst[i] == 4 else "y2a'1"
            rows.append(row)
        out = pd.DataFrame(rows)
        for key in ("source_neuron_idx", "target_neuron_idx"):
            out[key] = out[key].astype(np.int32)
        for key in ("source_body_id", "target_body_id"):
            out[key] = out[key].astype(np.uint64)
        # Match the observed frozen parquet object dtype; local pandas 3 may
        # infer its new string extension dtype from constructor literals.
        for prefix in ("source", "target"):
            for suffix in ("annotation_type", "annotation_instance"):
                out[prefix + "_" + suffix] = out[prefix + "_" + suffix].astype(object)
        out["edge_id"] = out.edge_id.astype(np.int64)
        out["raw_synapse_count"] = out.raw_synapse_count.astype(np.int64)
        return out
    candidates = table(range(4))
    candidates["plasticity_polarity"] = np.int64(-1)
    confirmed = table([4, 5])
    confirmed["source_compartment_verbatim"] = confirmed.target_compartment_verbatim
    confirmed["exact_compartment_match"] = True
    return meta, dan, mbon, confirmed, candidates, src, dst, raw


class Biology(unittest.TestCase):
    def test_complete_graph_binds_every_edge_array_length(self):
        meta, *_, src, dst, raw = fixture()
        graph = dict(nodes=np.arange(7, dtype=np.int32), edge_id=np.arange(6, dtype=np.int64),
                     src=src, dst=dst, signed_raw=raw)
        b._complete_graph_identity(meta, graph, n=7, e=6)
        for names in (("src",), ("dst",), ("signed_raw",), ("src", "dst", "signed_raw")):
            changed = dict(graph)
            for name in names:
                changed[name] = changed[name][:-1]
            with self.subTest(names=names), self.assertRaisesRegex(ValueError, "endpoint/weight lengths"):
                b._complete_graph_identity(meta, changed, n=7, e=6)
        for name in ("nodes", "edge_id"):
            changed = dict(graph)
            changed[name] = changed[name].reshape(1, -1)
            with self.subTest(name=name), self.assertRaises(ValueError):
                b._complete_graph_identity(meta, changed, n=7, e=6)

    def test_exact_id_compartment_and_real_edge_contract(self):
        values = fixture()
        populations, audit = b.audit_interface_tables(*values)
        self.assertEqual(populations["PAM15"].tolist(), [2])
        self.assertEqual(populations["MBON12"].tolist(), [5])
        self.assertEqual(audit["whole_DAN_mapping_unknown"], 1)
        self.assertFalse(audit["global_unknowns_filled"])
        self.assertEqual(audit["all_candidate_edges_retained"], 4)
        self.assertEqual(audit["interfaces"]["PAM15"]["candidate_contacts"], 16)
        self.assertEqual(audit["interfaces"]["PPL103"]["confirmed_DAN_MBON_contacts"], 24)
        self.assertGreater(audit["interfaces"]["PAM15"]["DAN_members"][0]["body_id"], 2**63)
        self.assertTrue(pd.isna(values[1].iloc[-1].mb_compartment_verbatim))

    def test_unknown_selected_mapping_is_not_filled_from_type(self):
        for field, new_value in (("mb_compartment_verbatim", None),
                                 ("mb_compartment_mapping_status", "unknown")):
            values = fixture()
            values[1].loc[2, field] = new_value
            with self.assertRaisesRegex(ValueError, "Unknown or mismatched"):
                b.audit_interface_tables(*values)

    def test_changed_body_id_or_type_or_instance_rejected(self):
        for field, new_value in (("body_id", np.uint64(5)), ("annotation_type", "PAM_OTHER"),
                                 ("annotation_instance", "inferred")):
            values = fixture()
            values[1].loc[2, field] = new_value
            with self.assertRaisesRegex(ValueError, "Identity field mismatch"):
                b.audit_interface_tables(*values)

    def test_unmapped_dan_cannot_be_silently_removed(self):
        values = list(fixture())
        values[1] = values[1].iloc[:2].copy()
        with self.assertRaisesRegex(ValueError, "complete exact-annotation"):
            b.audit_interface_tables(*values)

    def test_missing_actual_edge_fails_not_guessed(self):
        values = list(fixture())
        values[3] = values[3].iloc[:1].copy()
        with self.assertRaisesRegex(ValueError, "unknown or mismatched actual target"):
            b.audit_interface_tables(*values)

    def test_edge_identity_and_contacts_and_polarity(self):
        changes = ((3, "source_body_id", np.uint64(8)),
                   (3, "target_annotation_type", "MBON_OTHER"),
                   (4, "raw_synapse_count", np.int64(1)),
                   (4, "plasticity_polarity", np.int64(1)))
        for which, field, value in changes:
            values = fixture()
            values[which].loc[values[which].index[0], field] = value
            with self.subTest(field=field), self.assertRaises(ValueError):
                b.audit_interface_tables(*values)

    def test_false_or_nonboolean_confirmed_flag_rejected(self):
        for value in (False, "True"):
            values = fixture()
            values[3]["exact_compartment_match"] = value
            with self.assertRaisesRegex(ValueError, "unknown compartment"):
                b.audit_interface_tables(*values)

    def test_graph_endpoint_and_sign_change_rejected(self):
        for where in (5, 6, 7):
            values = fixture()
            values[where][0] = -1 if where == 7 else 6
            with self.assertRaises(ValueError):
                b.audit_interface_tables(*values)

    def test_schema_and_duplicate_identity_rejected(self):
        values = list(fixture())
        values[0] = values[0].drop(columns="annotation_class")
        with self.assertRaises(ValueError):
            b.audit_interface_tables(*values)
        values = fixture()
        values[0].loc[1, "body_id"] = values[0].loc[0, "body_id"]
        with self.assertRaises(ValueError):
            b.audit_interface_tables(*values)


if __name__ == "__main__":
    unittest.main()