#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EDA for the Konooz cross-domain Arabic NER data (KnowledgeGraphEval 2026, Subtask 1).

Parses the CoNLL-style domain files (one file per domain: ``<Domain>.txt``), where
each non-blank line is a token followed by N whitespace-separated tag columns (one
BIO channel per entity type in the *nested* representation), and blank lines mark
sentence boundaries.

The script is intentionally a plain, reproducible ``.py`` (not a notebook). It:
  * reports per-domain and corpus-level structural statistics (sentences, tokens,
    sentence-length distribution, type/token ratio, most frequent tokens);
  * confirms the tag-column schema and reports label density;
  * if gold labels are present (any non-O tag), decodes entities per type channel
    and reports entity counts, per-type distribution, and nesting statistics;
    if the data is unlabeled (all-O placeholders, i.e. the evaluation input), it
    says so explicitly and skips entity-level stats;
  * writes CSV summaries and (optionally) PNG plots to the output directory.

Usage
-----
    python eda_konooz.py \
        --input-dir ../../data/raw/konooz_dev \
        --out-dir   ../../data/eda/konooz_dev

External libraries (versions printed at runtime for replicability):
  * pandas          https://pandas.pydata.org/
  * matplotlib      https://matplotlib.org/   (optional; --no-plots to skip)
Standard library only otherwise. No network access.
"""
from __future__ import annotations

import argparse
import collections
import glob
import io
import os
import platform
import sys
from dataclasses import dataclass, field
from typing import Dict, List, Tuple

import pandas as pd

RANDOM_TOKENS_PREVIEW = 15


@dataclass
class DomainStats:
    domain: str
    n_sentences: int = 0
    n_tokens: int = 0
    n_tag_columns: int = 0
    sent_lengths: List[int] = field(default_factory=list)
    token_counter: "collections.Counter[str]" = field(default_factory=collections.Counter)
    non_o_tags: int = 0                       # (token, column) cells that are not "O"
    entities_by_type: "collections.Counter[str]" = field(default_factory=collections.Counter)
    nested_tokens: int = 0                    # tokens with >1 entity type active

    @property
    def n_types(self) -> int:
        return len(self.token_counter)

    @property
    def avg_sent_len(self) -> float:
        return (sum(self.sent_lengths) / len(self.sent_lengths)) if self.sent_lengths else 0.0

    @property
    def ttr(self) -> float:
        return (self.n_types / self.n_tokens) if self.n_tokens else 0.0


def read_conll(path: str) -> List[List[List[str]]]:
    """Return a list of sentences; each sentence is a list of ``[token, *tags]`` rows."""
    sentences: List[List[List[str]]] = []
    cur: List[List[str]] = []
    with io.open(path, encoding="utf-8") as fh:
        for line in fh:
            s = line.rstrip("\n")
            if not s.strip():
                if cur:
                    sentences.append(cur)
                    cur = []
                continue
            cur.append(s.split())
    if cur:
        sentences.append(cur)
    return sentences


def decode_spans(bio: List[str]) -> List[Tuple[int, int]]:
    """Decode (start, end) spans from a single B/I/O channel (lenient: I without B opens)."""
    spans, start = [], None
    for i, t in enumerate(list(bio) + ["O"]):
        tag = t.split("-", 1)[0] if t != "O" else "O"
        if tag == "B":
            if start is not None:
                spans.append((start, i - 1))
            start = i
        elif tag == "I":
            if start is None:
                start = i
        else:
            if start is not None:
                spans.append((start, i - 1))
                start = None
    return spans


def analyse_domain(domain: str, sentences: List[List[List[str]]]) -> DomainStats:
    st = DomainStats(domain=domain)
    st.n_sentences = len(sentences)
    # infer number of tag columns from the widest row
    st.n_tag_columns = max((len(r) - 1 for s in sentences for r in s), default=0)

    # per type-channel BIO sequences for entity decoding
    for sent in sentences:
        st.sent_lengths.append(len(sent))
        # columns: build a matrix of tags [token][col]
        col_seqs: Dict[int, List[str]] = collections.defaultdict(list)
        for row in sent:
            tok, tags = row[0], row[1:]
            st.n_tokens += 1
            st.token_counter[tok] += 1
            active = 0
            for c, t in enumerate(tags):
                col_seqs[c].append(t)
                if t != "O":
                    st.non_o_tags += 1
                    if t.split("-", 1)[0] in ("B",):
                        active += 1
            if active > 1:
                st.nested_tokens += 1
        # decode entities per column; a column header/type name is unknown here,
        # so we label entities by their tag suffix (type) when present, else col index.
        for c, seq in col_seqs.items():
            for (s0, s1) in decode_spans(seq):
                # type = suffix of the B- tag if present
                btag = seq[s0]
                typ = btag.split("-", 1)[1] if "-" in btag else f"col{c}"
                st.entities_by_type[typ] += 1
    return st


def summarise(all_stats: List[DomainStats]) -> pd.DataFrame:
    rows = []
    for st in all_stats:
        rows.append({
            "domain": st.domain,
            "sentences": st.n_sentences,
            "tokens": st.n_tokens,
            "types": st.n_types,
            "type_token_ratio": round(st.ttr, 4),
            "avg_sent_len": round(st.avg_sent_len, 2),
            "max_sent_len": max(st.sent_lengths) if st.sent_lengths else 0,
            "tag_columns": st.n_tag_columns,
            "non_O_cells": st.non_o_tags,
            "entities": sum(st.entities_by_type.values()),
            "nested_tokens": st.nested_tokens,
        })
    df = pd.DataFrame(rows).sort_values("domain").reset_index(drop=True)
    total = {
        "domain": "ALL",
        "sentences": df["sentences"].sum(),
        "tokens": df["tokens"].sum(),
        "types": None,
        "type_token_ratio": None,
        "avg_sent_len": round(df["tokens"].sum() / max(df["sentences"].sum(), 1), 2),
        "max_sent_len": df["max_sent_len"].max(),
        "tag_columns": df["tag_columns"].max(),
        "non_O_cells": df["non_O_cells"].sum(),
        "entities": df["entities"].sum(),
        "nested_tokens": df["nested_tokens"].sum(),
    }
    return pd.concat([df, pd.DataFrame([total])], ignore_index=True)


def maybe_plot(df: pd.DataFrame, out_dir: str) -> None:
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except Exception as exc:  # pragma: no cover
        print(f"[plot] matplotlib unavailable ({exc}); skipping plots.")
        return
    body = df[df["domain"] != "ALL"]
    # tokens per domain
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.bar(body["domain"], body["tokens"])
    ax.set_ylabel("tokens"); ax.set_title("Konooz dev: tokens per domain")
    ax.tick_params(axis="x", rotation=45)
    fig.tight_layout(); fig.savefig(os.path.join(out_dir, "konooz_tokens_per_domain.png"), dpi=150)
    plt.close(fig)
    # sentences per domain
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.bar(body["domain"], body["sentences"])
    ax.set_ylabel("sentences"); ax.set_title("Konooz dev: sentences per domain")
    ax.tick_params(axis="x", rotation=45)
    fig.tight_layout(); fig.savefig(os.path.join(out_dir, "konooz_sentences_per_domain.png"), dpi=150)
    plt.close(fig)
    print(f"[plot] wrote 2 PNGs to {out_dir}")


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    here = os.path.dirname(os.path.abspath(__file__))
    ap.add_argument("--input-dir", default=os.path.join(here, "..", "..", "data", "raw", "konooz_dev"),
                    help="directory of <Domain>.txt CoNLL files")
    ap.add_argument("--out-dir", default=os.path.join(here, "..", "..", "data", "eda", "konooz_dev"),
                    help="where to write CSV summaries and plots")
    ap.add_argument("--no-plots", action="store_true", help="skip PNG plots")
    args = ap.parse_args(argv)

    print(f"[env] python {platform.python_version()} | pandas {pd.__version__}")
    in_dir = os.path.abspath(args.input_dir)
    out_dir = os.path.abspath(args.out_dir)
    os.makedirs(out_dir, exist_ok=True)

    files = sorted(glob.glob(os.path.join(in_dir, "*.txt")))
    if not files:
        print(f"[error] no .txt files under {in_dir}", file=sys.stderr)
        return 1
    print(f"[data] {len(files)} domain files under {in_dir}")

    all_stats: List[DomainStats] = []
    for f in files:
        domain = os.path.splitext(os.path.basename(f))[0]
        all_stats.append(analyse_domain(domain, read_conll(f)))

    summary = summarise(all_stats)
    summary_path = os.path.join(out_dir, "konooz_dev_summary.csv")
    summary.to_csv(summary_path, index=False, encoding="utf-8")

    total_non_o = int(summary.loc[summary["domain"] == "ALL", "non_O_cells"].iloc[0])
    labelled = total_non_o > 0

    print("\n===== Konooz dev EDA =====")
    with pd.option_context("display.max_columns", None, "display.width", 200):
        print(summary.to_string(index=False))
    n_cols = int(summary.loc[summary["domain"] == "ALL", "tag_columns"].iloc[0])
    print(f"\n[schema] token + {n_cols} tag columns per line "
          f"(nested representation: one B/I/O channel per entity type).")
    if not labelled:
        print("[labels] ALL tag cells are 'O' -> this is the UNLABELED evaluation "
              "input (participants predict labels). Entity-level stats skipped.")
    else:
        print(f"[labels] {total_non_o} non-O cells found -> decoding entity statistics.")
        ent = collections.Counter()
        nested = 0
        for st in all_stats:
            ent.update(st.entities_by_type)
            nested += st.nested_tokens
        ent_df = (pd.DataFrame(sorted(ent.items(), key=lambda kv: -kv[1]),
                               columns=["entity_type", "count"]))
        ent_path = os.path.join(out_dir, "konooz_dev_entity_types.csv")
        ent_df.to_csv(ent_path, index=False, encoding="utf-8")
        print(ent_df.to_string(index=False))
        print(f"[nesting] tokens with >1 active entity type: {nested}")
        print(f"[write] {ent_path}")

    # top tokens overall (lexical signal / OOV inspection)
    overall_tokens = collections.Counter()
    for st in all_stats:
        overall_tokens.update(st.token_counter)
    top = overall_tokens.most_common(50)
    pd.DataFrame(top, columns=["token", "count"]).to_csv(
        os.path.join(out_dir, "konooz_dev_top_tokens.csv"), index=False, encoding="utf-8")
    preview = ", ".join(t for t, _ in top[:RANDOM_TOKENS_PREVIEW])
    print(f"\n[lexicon] {len(overall_tokens):,} unique tokens; most frequent: {preview}")

    if not args.no_plots:
        maybe_plot(summary, out_dir)

    print(f"\n[write] {summary_path}")
    print("[done]")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())