#!/usr/bin/env python3
"""Evaluate spatial transcriptomics analysis output against reference.

Checks multiple output files for validity and scores clustering against
pathologist-annotated cortical layer labels using Adjusted Rand Index.

Required outputs (in output/):
  qc_stats.txt          - QC filtering summary
  cell_clusters.csv     - Spot-level cluster assignments
  marker_genes.csv      - Top marker genes per cluster
  umap_plot.png         - UMAP visualization
  results.h5ad          - Processed AnnData

Scoring weights:
  qc_stats.txt          15%
  cell_clusters.csv     15%
  marker_genes.csv      10%
  umap_plot.png          5%
  results.h5ad          15%
  Clustering ARI        40%  (against pathologist annotations)
"""

import csv
import json
import os
import struct
import sys
import zlib
from collections import Counter


def _load_input_barcodes(input_dir):
    """Load the set of valid in-tissue barcodes from the input data."""
    barcodes = set()
    # Try tissue_positions_list.csv or .txt
    for name in ['tissue_positions_list.csv', 'tissue_positions_list.txt']:
        path = os.path.join(input_dir, 'spatial', name)
        if os.path.exists(path):
            with open(path) as f:
                for line in f:
                    parts = line.strip().split(',')
                    if len(parts) >= 2 and parts[1] == '1':  # in_tissue=1
                        barcodes.add(parts[0])
            break
    # Fallback: try h5 matrix barcodes
    if not barcodes:
        try:
            import h5py
            h5_path = os.path.join(input_dir, 'filtered_feature_bc_matrix.h5')
            if os.path.exists(h5_path):
                with h5py.File(h5_path, 'r') as f:
                    for grp in ['matrix', '/']:
                        if 'barcodes' in f[grp]:
                            barcodes = {b.decode() for b in f[grp]['barcodes'][:]}
                            break
        except (ImportError, Exception):
            pass
    return barcodes


def check_png(path):
    """Check if file is a valid PNG with full decode if possible."""
    if not os.path.exists(path):
        return 0.0, "file not found"
    size = os.path.getsize(path)
    if size < 2000:
        return 0.0, "file too small (%d bytes, need > 2000)" % size
    with open(path, 'rb') as f:
        header = f.read(8)
    if header != b'\x89PNG\r\n\x1a\n':
        return 0.0, "not a valid PNG (bad header)"

    # Try full decode with PIL
    try:
        from PIL import Image
        img = Image.open(path)
        img.verify()
        img = Image.open(path)  # re-open after verify
        w, h = img.size
        if w < 100 or h < 100:
            return 0.5, "valid PNG but very small (%dx%d)" % (w, h)
        return 1.0, "valid PNG, %dx%d, %d bytes" % (w, h, size)
    except ImportError:
        return 0.8, "valid PNG header, %d bytes (PIL not available for full decode)" % size
    except Exception as e:
        return 0.0, "PNG decode failed: %s" % e


def check_qc_stats(path):
    """Check QC stats file."""
    if not os.path.exists(path):
        return 0.0, "qc_stats.txt not found"
    try:
        with open(path) as f:
            content = f.read()
        lines = content.strip().split('\n')
        stats = {}
        for line in lines:
            for sep in [':', '=']:
                if sep in line:
                    key, val = line.split(sep, 1)
                    stats[key.strip().lower().replace(' ', '_')] = val.strip()
                    break

        retained = None
        for key in ['spots_after', 'spots_retained', 'retained_spots',
                     'n_spots_after', 'cells_retained', 'n_obs',
                     'spots_after_qc', 'cells_after']:
            if key in stats:
                try:
                    retained = int(stats[key].replace(',', ''))
                    break
                except ValueError:
                    continue

        if retained is None:
            return 0.3, "parsed but spots_after not found. Keys: %s" % list(stats.keys())

        # Reasonable range: at least 2000 spots retained from ~3639
        if 2000 <= retained <= 3700:
            return 1.0, "spots_after=%d (reasonable)" % retained
        else:
            return 0.3, "spots_after=%d (outside 2000-3700)" % retained
    except Exception as e:
        return 0.0, "parse failed: %s" % e


def check_clusters(path, input_barcodes=None):
    """Check cell_clusters.csv with barcode validation."""
    if not os.path.exists(path):
        return 0.0, [], "cell_clusters.csv not found"
    try:
        with open(path) as f:
            reader = csv.DictReader(f)
            rows = list(reader)

        if not rows:
            return 0.0, [], "empty file"

        cols = set(rows[0].keys())
        bcol = next((c for c in ['barcode', 'Barcode'] if c in cols), None)
        ccol = next((c for c in ['cluster', 'Cluster', 'leiden'] if c in cols), None)

        if not bcol:
            return 0.1, [], "missing barcode column. Found: %s" % cols
        if not ccol:
            return 0.1, [], "missing cluster column. Found: %s" % cols

        all_barcodes = [r[bcol] for r in rows]
        unique_barcodes = set(all_barcodes)
        clusters = set(r[ccol] for r in rows)

        warnings = []
        score = 1.0

        # Check for duplicates
        if len(all_barcodes) != len(unique_barcodes):
            n_dup = len(all_barcodes) - len(unique_barcodes)
            score *= 0.7
            warnings.append("%d duplicate barcodes" % n_dup)

        # Validate barcodes against input
        if input_barcodes:
            invalid = unique_barcodes - input_barcodes
            if invalid:
                frac_invalid = len(invalid) / len(unique_barcodes)
                if frac_invalid > 0.1:
                    score *= 0.5
                warnings.append("%d/%d barcodes not in input" % (
                    len(invalid), len(unique_barcodes)))

        if not (2000 <= len(unique_barcodes) <= 3700):
            score *= 0.5
            warnings.append("spots outside 2000-3700")
        if not (4 <= len(clusters) <= 15):
            score *= 0.5
            warnings.append("clusters outside 4-15")

        detail = "n_spots=%d, n_clusters=%d" % (len(unique_barcodes), len(clusters))
        if warnings:
            detail += " [%s]" % "; ".join(warnings)

        return score, list(clusters), detail
    except Exception as e:
        return 0.0, [], "parse failed: %s" % e


def check_markers(path, expected_clusters):
    """Check marker_genes.csv with column and numeric validation."""
    if not os.path.exists(path):
        return 0.0, "marker_genes.csv not found"
    try:
        with open(path) as f:
            reader = csv.DictReader(f)
            rows = list(reader)

        if not rows:
            return 0.0, "empty file"

        cols = set(rows[0].keys())
        ccol = next((c for c in cols if 'cluster' in c.lower()), None)
        gcol = next((c for c in cols if 'gene' in c.lower() or 'name' in c.lower()), None)

        if not ccol or not gcol:
            return 0.1, "missing cluster or gene column. Found: %s" % cols

        # Check for score and pval_adj columns
        scol = next((c for c in cols if c.lower() in [
            'score', 'scores', 'logfoldchange', 'logfoldchanges',
            'log2fc', 'avg_log2fc']), None)
        pcol = next((c for c in cols if 'pval' in c.lower() or 'padj' in c.lower()
                     or 'fdr' in c.lower() or 'p_val' in c.lower()), None)

        col_score = 1.0
        col_warnings = []

        if not scol:
            col_score *= 0.7
            col_warnings.append("no score/logfoldchange column")
        if not pcol:
            col_score *= 0.7
            col_warnings.append("no pval_adj column")

        # Validate numeric values where columns exist
        n_invalid_score = 0
        n_invalid_pval = 0
        for r in rows:
            if scol and r.get(scol):
                try:
                    float(r[scol])
                except (ValueError, TypeError):
                    n_invalid_score += 1
            if pcol and r.get(pcol):
                try:
                    v = float(r[pcol])
                    if v < 0 or v > 1:
                        n_invalid_pval += 1
                except (ValueError, TypeError):
                    n_invalid_pval += 1

        if n_invalid_score > 0:
            col_score *= 0.8
            col_warnings.append("%d non-numeric scores" % n_invalid_score)
        if n_invalid_pval > 0:
            col_score *= 0.8
            col_warnings.append("%d invalid p-values" % n_invalid_pval)

        by_cluster = {}
        for r in rows:
            c = str(r[ccol])
            by_cluster.setdefault(c, []).append(r)

        if not expected_clusters:
            detail = "markers for %d clusters" % len(by_cluster)
            if col_warnings:
                detail += " [%s]" % "; ".join(col_warnings)
            return 0.5 * col_score, detail

        covered = sum(1 for ec in expected_clusters
                      if len(by_cluster.get(str(ec), [])) >= 5)
        coverage = covered / len(expected_clusters) if expected_clusters else 0
        detail = "%d/%d clusters with >=5 markers" % (covered, len(expected_clusters))
        if col_warnings:
            detail += " [%s]" % "; ".join(col_warnings)
        return coverage * col_score, detail
    except Exception as e:
        return 0.0, "parse failed: %s" % e


def check_h5ad(path):
    """Check results.h5ad."""
    if not os.path.exists(path):
        return 0.0, "results.h5ad not found"

    size = os.path.getsize(path)
    if size < 1000:
        return 0.0, "file too small (%d bytes)" % size

    try:
        import anndata
        adata = anndata.read_h5ad(path)

        checks = []
        score = 0.0

        if 'X_pca' in adata.obsm:
            score += 0.33
            checks.append("X_pca ok")
        else:
            checks.append("X_pca MISSING")

        if 'X_umap' in adata.obsm:
            score += 0.33
            checks.append("X_umap ok")
        else:
            checks.append("X_umap MISSING")

        has_cluster = ('leiden' in adata.obs.columns or
                       'cluster' in adata.obs.columns or
                       'louvain' in adata.obs.columns)
        if has_cluster:
            score += 0.34
            checks.append("clustering ok")
        else:
            checks.append("clustering MISSING")

        return score, "h5ad (%d obs, %d vars). %s" % (
            adata.n_obs, adata.n_vars, ", ".join(checks))
    except ImportError:
        try:
            import h5py
            with h5py.File(path, 'r') as f:
                keys = list(f.keys())
            if 'obs' in keys and 'var' in keys:
                return 0.5, "valid HDF5 with obs/var (anndata not available)"
            return 0.2, "valid HDF5 but missing obs/var"
        except Exception:
            return 0.0, "not a valid HDF5"
    except Exception as e:
        return 0.0, "load failed: %s" % e


def comb2(n):
    """Binomial coefficient C(n, 2)."""
    return n * (n - 1) / 2.0


def adjusted_rand_index(labels_true, labels_pred):
    """Compute ARI without sklearn."""
    n = len(labels_true)
    if n == 0:
        return 0.0

    contingency = Counter(zip(labels_true, labels_pred))
    row_sums = Counter(labels_true)
    col_sums = Counter(labels_pred)

    sum_comb_c = sum(comb2(v) for v in contingency.values())
    sum_comb_a = sum(comb2(v) for v in row_sums.values())
    sum_comb_b = sum(comb2(v) for v in col_sums.values())
    comb_n = comb2(n)

    if comb_n == 0:
        return 0.0

    expected = sum_comb_a * sum_comb_b / comb_n
    max_index = 0.5 * (sum_comb_a + sum_comb_b)

    if max_index == expected:
        return 1.0 if sum_comb_c == expected else 0.0

    return (sum_comb_c - expected) / (max_index - expected)


def compute_ari(pred_path, ref_path, input_barcodes=None):
    """Compute ARI between predicted clusters and reference labels.

    Validates predicted barcodes against the input barcode set and
    requires near-complete overlap with reference labels.
    """
    if not os.path.exists(pred_path):
        return 0.0, "cell_clusters.csv not found"
    if not os.path.exists(ref_path):
        return 0.0, "reference_labels.csv not found"

    try:
        pred_labels = {}
        with open(pred_path) as f:
            reader = csv.DictReader(f)
            cols = reader.fieldnames
            bcol = next((c for c in ['barcode', 'Barcode'] if c in cols), None)
            ccol = next((c for c in ['cluster', 'Cluster', 'leiden'] if c in cols), None)
            if not bcol or not ccol:
                return 0.0, "missing barcode/cluster columns in predictions"
            for r in reader:
                bc = r[bcol]
                if bc not in pred_labels:  # take first occurrence if duplicates
                    pred_labels[bc] = r[ccol]

        ref_labels = {}
        with open(ref_path) as f:
            reader = csv.DictReader(f)
            cols = reader.fieldnames
            bcol = next((c for c in ['barcode', 'Barcode'] if c in cols), None)
            lcol = next((c for c in ['label', 'cluster', 'domain', 'layer'] if c in cols), None)
            if not bcol or not lcol:
                return 0.0, "missing barcode/label columns in reference"
            for r in reader:
                ref_labels[r[bcol]] = r[lcol]

        # Validate predicted barcodes against input if available
        warnings = []
        if input_barcodes:
            invalid = set(pred_labels.keys()) - input_barcodes
            if invalid:
                warnings.append("%d predicted barcodes not in input" % len(invalid))

        common = set(pred_labels.keys()) & set(ref_labels.keys())
        n_ref = len(ref_labels)

        # Require substantial overlap: at least 50% of reference barcodes
        min_overlap = max(n_ref // 2, 500)
        if len(common) < min_overlap:
            return 0.0, "insufficient overlap (%d < %d, ref has %d)" % (
                len(common), min_overlap, n_ref)

        # Apply coverage penalty if overlap is below 90% of reference
        coverage = len(common) / n_ref if n_ref > 0 else 0
        coverage_penalty = 1.0
        if coverage < 0.90:
            coverage_penalty = coverage / 0.90

        pred = [pred_labels[b] for b in common]
        ref = [ref_labels[b] for b in common]
        ari = adjusted_rand_index(ref, pred)

        # Score: full credit at ARI >= 0.50, partial at >= 0.25
        if ari >= 0.50:
            score = 1.0
        elif ari >= 0.25:
            score = (ari - 0.25) / 0.25  # linear 0→1
        else:
            score = 0.0

        score *= coverage_penalty

        detail = "ARI=%.4f (overlap=%d/%d=%.1f%%, coverage_penalty=%.2f, score=%.2f)" % (
            ari, len(common), n_ref, 100 * coverage, coverage_penalty, score)
        if warnings:
            detail += " [%s]" % "; ".join(warnings)
        return score, detail
    except Exception as e:
        return 0.0, "ARI failed: %s" % e


def evaluate():
    """Run full evaluation."""
    if len(sys.argv) >= 4:
        input_dir = sys.argv[1]
        output_dir = sys.argv[2]
        ref_dir = sys.argv[3]
    else:
        input_dir = 'input'
        output_dir = 'output'
        ref_dir = os.path.join(
            os.path.dirname(os.path.abspath(__file__)), '..', 'reference')
        if not os.path.isdir(ref_dir):
            ref_dir = 'reference'

    # Load input barcodes for validation
    input_barcodes = _load_input_barcodes(input_dir)

    results = {}
    total = 0.0

    # 1. QC stats (15%)
    s, d = check_qc_stats(os.path.join(output_dir, 'qc_stats.txt'))
    results['qc_stats'] = {'score': s, 'weight': 0.15, 'detail': d}
    total += 0.15 * s

    # 2. Clusters (15%)
    s, cl, d = check_clusters(
        os.path.join(output_dir, 'cell_clusters.csv'), input_barcodes)
    results['cell_clusters'] = {'score': s, 'weight': 0.15, 'detail': d}
    total += 0.15 * s

    # 3. Markers (10%)
    s, d = check_markers(os.path.join(output_dir, 'marker_genes.csv'), cl)
    results['marker_genes'] = {'score': s, 'weight': 0.10, 'detail': d}
    total += 0.10 * s

    # 4. UMAP plot (5%)
    s, d = check_png(os.path.join(output_dir, 'umap_plot.png'))
    results['umap_plot'] = {'score': s, 'weight': 0.05, 'detail': d}
    total += 0.05 * s

    # 5. results.h5ad (15%)
    s, d = check_h5ad(os.path.join(output_dir, 'results.h5ad'))
    results['results_h5ad'] = {'score': s, 'weight': 0.15, 'detail': d}
    total += 0.15 * s

    # 6. ARI (40%)
    s, d = compute_ari(
        os.path.join(output_dir, 'cell_clusters.csv'),
        os.path.join(ref_dir, 'reference_labels.csv'),
        input_barcodes)
    results['clustering_ari'] = {'score': s, 'weight': 0.40, 'detail': d}
    total += 0.40 * s

    results['total'] = round(total, 4)

    print("SCORE: %.4f" % total)
    print(json.dumps(results, indent=2))


if __name__ == '__main__':
    evaluate()