"""Organization normalization and team bloc detection."""

from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass, field

from .db import Database

# Maps raw affiliation strings to canonical org names.
# Built from SELECT DISTINCT affiliation FROM draft_authors.
ORG_ALIASES: dict[str, str] = {
    # Huawei
    "Huawei Technologies": "Huawei",
    "Huawei Technologies Co., Ltd.": "Huawei",
    "Huawei Technologies, Co., Ltd": "Huawei",
    "Huawei Tech": "Huawei",
    "Huawei Canada": "Huawei",
    "Huawei R&D": "Huawei",
    "Huawei Singapore": "Huawei",
    # Cisco
    "Cisco Systems": "Cisco",
    "Cisco Systems, Inc.": "Cisco",
    # Ericsson
    "Ericsson AB": "Ericsson",
    # RISE
    "RISE AB": "RISE",
    "RISE": "RISE",
    # Independent
    "Independent Researcher": "Independent",
    "Unaffiliated": "Independent",
    "Individual Contributor": "Independent",
    # Inria
    "INRIA": "Inria",
    # Google
    "Google LLC": "Google",
    "Google": "Google",
    # Apple
    "Apple Inc": "Apple",
    "Apple, Inc": "Apple",
    "Apple": "Apple",
    # Amazon
    "Amazon Web Services": "Amazon",
    "AWS": "Amazon",
    "Amazon": "Amazon",
    # Siemens
    "Siemens AG": "Siemens",
    # ZTE
    "ZTE": "ZTE Corporation",
    # Telefonica
    "Telefonica I+D": "Telefonica",
    # Deutsche Telekom
    "Deutsche Telecom": "Deutsche Telekom",
    # InterDigital
    "InterDigital Europe Ltd.": "InterDigital Europe",
    # Boeing
    "Boeing Technology Innovation": "Boeing",
    "Boeing Research & Technology": "Boeing",
    # Futurewei
    "Futurewei Technologies USA": "Futurewei",
    "Futurewei": "Futurewei",
    # IBM
    "IBM Research": "IBM",
    "IBM": "IBM",
    # China Telecom
    "China Telecom Research Institute": "China Telecom",
    # Beijing University (multiline variant from Datatracker)
    "Beijing University of Posts and\n      Telecommunications": "BUPT",
    "Beijing University of Posts and Telecommunications": "BUPT",
    # AsiaInfo
    "AsiaInfo Technologies (China) Inc.": "AsiaInfo",
    "AsiaInfo Technologies (China) Inc": "AsiaInfo",
    # Dept of CS
    "Department of Computer Science and Engineering": "Department of Computer Science & Engineering",
}

# Common suffixes to strip for fuzzy matching
_SUFFIXES = [
    ", Inc.", ", Inc", " Inc.", " Inc",
    " LLC", " Ltd.", " Ltd",
    " AB", " GmbH", " Corp",
    " Co., Ltd.", " Co., Ltd",
    " Technologies",
]


def normalize_org(raw: str) -> str:
    """Normalize an affiliation string to a canonical org name."""
    raw = raw.strip()
    if not raw:
        return ""
    # Exact match
    if raw in ORG_ALIASES:
        return ORG_ALIASES[raw]
    # Fuzzy: strip suffixes and check again
    stripped = raw
    for suffix in _SUFFIXES:
        if stripped.endswith(suffix):
            stripped = stripped[: -len(suffix)].strip()
            break
    if stripped in ORG_ALIASES:
        return ORG_ALIASES[stripped]
    if stripped != raw and stripped:
        # Check if the stripped form matches a canonical name directly
        for canonical in set(ORG_ALIASES.values()):
            if stripped.lower() == canonical.lower():
                return canonical
    return raw


@dataclass
class Bloc:
    """A team of authors who consistently co-author together."""

    members: list[tuple[int, str, str]]  # (person_id, name, normalized_org)
    shared_drafts: int  # drafts where >= 2 members co-author
    primary_org: str
    cohesion: float  # avg pairwise cohesion

    @property
    def member_pids(self) -> set[int]:
        return {pid for pid, _, _ in self.members}

    @property
    def label(self) -> str:
        return f"{self.primary_org} team ({len(self.members)})"


def detect_blocs(
    db: Database,
    cohesion_threshold: float = 0.70,
    min_size: int = 2,
    min_shared_drafts: int = 2,
) -> list[Bloc]:
    """Detect team blocs where all member pairs share >= threshold of drafts.

    Uses connected components on a cohesion-filtered co-author graph,
    then merges overlapping groups into single blocs.
    """
    draft_counts = db.author_draft_counts()
    draft_sets = db.author_draft_sets()

    # Get enriched pair data with person_ids
    rows = db.conn.execute(
        """SELECT a1.name, da1.person_id, a2.name, da2.person_id, COUNT(*) as shared
        FROM draft_authors da1
        JOIN draft_authors da2 ON da1.draft_name = da2.draft_name
            AND da1.person_id < da2.person_id
        JOIN authors a1 ON da1.person_id = a1.person_id
        JOIN authors a2 ON da2.person_id = a2.person_id
        GROUP BY da1.person_id, da2.person_id
        HAVING shared >= ?
        ORDER BY shared DESC""",
        (min_shared_drafts,),
    ).fetchall()

    # Get affiliations per person
    aff_rows = db.conn.execute(
        "SELECT person_id, affiliation FROM authors"
    ).fetchall()
    person_aff = {r[0]: normalize_org(r[1]) for r in aff_rows}
    person_name: dict[int, str] = {}

    # Build cohesion-filtered adjacency: only keep edges with high overlap
    adj: dict[int, set[int]] = defaultdict(set)
    pair_shared: dict[tuple[int, int], int] = {}
    pair_cohesion: dict[tuple[int, int], float] = {}

    for r in rows:
        name_a, pid_a, name_b, pid_b, shared = r[0], r[1], r[2], r[3], r[4]
        person_name[pid_a] = name_a
        person_name[pid_b] = name_b
        min_d = min(draft_counts.get(pid_a, 1), draft_counts.get(pid_b, 1))
        cohesion = shared / min_d
        if cohesion >= cohesion_threshold:
            adj[pid_a].add(pid_b)
            adj[pid_b].add(pid_a)
            key = (min(pid_a, pid_b), max(pid_a, pid_b))
            pair_shared[key] = shared
            pair_cohesion[key] = cohesion

    # Find connected components (each component = one merged bloc)
    visited: set[int] = set()
    components: list[set[int]] = []
    for pid in adj:
        if pid in visited:
            continue
        component: set[int] = set()
        stack = [pid]
        while stack:
            node = stack.pop()
            if node in visited:
                continue
            visited.add(node)
            component.add(node)
            stack.extend(adj[node] - visited)
        if len(component) >= min_size:
            components.append(component)

    # Build Bloc objects from components
    blocs = []
    for comp in components:
        members = [
            (pid, person_name.get(pid, "?"), person_aff.get(pid, ""))
            for pid in comp
        ]

        # Shared drafts = drafts where >= 2 bloc members appear
        all_drafts: dict[str, int] = defaultdict(int)
        for pid in comp:
            for d in draft_sets.get(pid, set()):
                all_drafts[d] += 1
        shared_count = sum(1 for cnt in all_drafts.values() if cnt >= 2)

        # Primary org = most common among members
        org_counts: dict[str, int] = defaultdict(int)
        for _, _, org in members:
            if org:
                org_counts[org] += 1
        primary = max(org_counts, key=org_counts.get) if org_counts else ""

        # Average pairwise cohesion (only for connected pairs)
        edges = [
            pair_cohesion[key]
            for a in comp for b in comp if a < b
            for key in [(a, b)] if key in pair_cohesion
        ]
        avg_coh = sum(edges) / len(edges) if edges else 0

        blocs.append(Bloc(
            members=sorted(members, key=lambda m: -len(draft_sets.get(m[0], set()))),
            shared_drafts=shared_count,
            primary_org=primary,
            cohesion=avg_coh,
        ))

    # Sort: most shared drafts first (the interesting ones)
    blocs.sort(key=lambda b: (-b.shared_drafts, -len(b.members)))
    return blocs


def top_orgs_normalized(
    db: Database, limit: int = 20
) -> list[tuple[str, int, int]]:
    """Return (canonical_org, unique_authors, unique_drafts) with merged orgs."""
    raw = db.org_data_raw()

    org_people: dict[str, set[int]] = defaultdict(set)
    org_drafts: dict[str, set[str]] = defaultdict(set)

    for aff, pid, draft_name in raw:
        canonical = normalize_org(aff)
        if canonical:
            org_people[canonical].add(pid)
            org_drafts[canonical].add(draft_name)

    results = [
        (org, len(org_people[org]), len(org_drafts[org]))
        for org in org_people
    ]
    results.sort(key=lambda x: -x[2])
    return results[:limit]


def cross_org_normalized(
    db: Database, limit: int = 20
) -> list[tuple[str, str, int]]:
    """Return (org_a, org_b, shared_drafts) with normalized org names."""
    # Get all (aff_a, aff_b, draft_name) cross-org triples
    rows = db.conn.execute(
        """SELECT da1.affiliation, da2.affiliation, da1.draft_name
        FROM draft_authors da1
        JOIN draft_authors da2 ON da1.draft_name = da2.draft_name
            AND da1.person_id < da2.person_id
        WHERE da1.affiliation != '' AND da2.affiliation != ''"""
    ).fetchall()

    pair_drafts: dict[tuple[str, str], set[str]] = defaultdict(set)
    for aff_a, aff_b, draft_name in rows:
        norm_a = normalize_org(aff_a)
        norm_b = normalize_org(aff_b)
        if norm_a and norm_b and norm_a != norm_b:
            key = tuple(sorted([norm_a, norm_b]))
            pair_drafts[key].add(draft_name)

    results = [
        (org_a, org_b, len(drafts))
        for (org_a, org_b), drafts in pair_drafts.items()
    ]
    results.sort(key=lambda x: -x[2])
    return results[:limit]
