#!/usr/bin/env python3
import argparse
import csv
import os
import re
import ssl
import statistics
import sys
import urllib.error
import urllib.parse
import urllib.request
import zipfile
import xml.etree.ElementTree as ET
from collections import defaultdict


def col_to_index(col):
    idx = 0
    for c in col:
        idx = idx * 26 + (ord(c.upper()) - ord("A") + 1)
    return idx


def parse_shared_strings(zf):
    try:
        xml = zf.read("xl/sharedStrings.xml")
    except KeyError:
        return []
    root = ET.fromstring(xml)
    ns = {"a": root.tag.split("}")[0].strip("{")}
    shared = []
    for si in root.findall(".//a:si", ns):
        text_parts = []
        for t in si.findall(".//a:t", ns):
            text_parts.append(t.text or "")
        shared.append("".join(text_parts))
    return shared


def list_sheets(zf):
    wb = ET.fromstring(zf.read("xl/workbook.xml"))
    ns = {"a": wb.tag.split("}")[0].strip("{")}
    sheets = []
    for sh in wb.findall(".//a:sheets/a:sheet", ns):
        sheets.append(
            (sh.attrib["name"], sh.attrib["{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id"])
        )

    rels = ET.fromstring(zf.read("xl/_rels/workbook.xml.rels"))
    relmap = {}
    for rel in rels.findall(".//{http://schemas.openxmlformats.org/package/2006/relationships}Relationship"):
        relmap[rel.attrib["Id"]] = rel.attrib["Target"]

    return [(name, "xl/" + relmap[rid]) for name, rid in sheets]


def read_xlsx_sheet(path, sheet_name):
    with zipfile.ZipFile(path) as zf:
        shared = parse_shared_strings(zf)
        sheets = dict(list_sheets(zf))
        if sheet_name not in sheets:
            raise ValueError(f"Sheet '{sheet_name}' not found. Available: {', '.join(sheets.keys())}")

        root = ET.fromstring(zf.read(sheets[sheet_name]))
        ns = {"a": root.tag.split("}")[0].strip("{")}
        row_dicts = []
        max_col_idx = 0

        for row in root.findall(".//a:sheetData/a:row", ns):
            row_cells = {}
            for c in row.findall("a:c", ns):
                ref = c.attrib.get("r", "")
                col = "".join(ch for ch in ref if ch.isalpha())
                if not col:
                    continue
                v = c.find("a:v", ns)
                value = v.text if v is not None else ""
                if c.attrib.get("t") == "s":
                    try:
                        value = shared[int(value)]
                    except Exception:
                        pass
                row_cells[col] = value
                max_col_idx = max(max_col_idx, col_to_index(col))
            if row_cells:
                row_dicts.append(row_cells)

        if not row_dicts:
            return []

        rows = []
        for row_cells in row_dicts:
            row_list = [""] * max_col_idx
            for col, val in row_cells.items():
                row_list[col_to_index(col) - 1] = val
            rows.append(row_list)

        header = rows[0]
        ncols = len(header)
        while ncols > 0 and header[ncols - 1] == "":
            ncols -= 1

        normalized = []
        for row in rows:
            if len(row) < ncols:
                row = row + [""] * (ncols - len(row))
            elif len(row) > ncols:
                row = row[:ncols]
            normalized.append(row)
        return normalized


def score_counts(values, metric):
    if not values:
        return 0.0
    if metric == "mean":
        return sum(values) / len(values)
    if metric == "median":
        return statistics.median(values)
    if metric == "sum":
        return sum(values)
    if metric == "max":
        return max(values)
    raise ValueError(f"Unknown metric: {metric}")


def canonical_mirna(name):
    s = name.strip().lower()
    s = re.sub(r"^[a-z]{3}-", "", s)
    return s


def _detect_delimiter(sample):
    if "\t" in sample:
        return "\t"
    if "," in sample:
        return ","
    return "\t"


def read_delimited_rows(path):
    with open(path, "r", newline="") as f:
        sample = f.readline()
        delim = _detect_delimiter(sample)
        f.seek(0)
        try:
            reader = csv.reader(f, delimiter=delim)
            rows = list(reader)
        except csv.Error:
            f.seek(0)
            reader = csv.reader(
                f,
                delimiter=delim,
                quoting=csv.QUOTE_NONE,
                escapechar="\\",
            )
            rows = list(reader)
    return rows


def iter_delimited_rows(path):
    with open(path, "r", newline="") as f:
        sample = f.readline()
        delim = _detect_delimiter(sample)
        f.seek(0)
        try:
            reader = csv.reader(f, delimiter=delim)
            for row in reader:
                yield row
        except csv.Error:
            f.seek(0)
            reader = csv.reader(
                f,
                delimiter=delim,
                quoting=csv.QUOTE_NONE,
                escapechar="\\",
            )
            for row in reader:
                yield row


def read_targets(path):
    rows = read_delimited_rows(path)
    return parse_targets_rows(rows)


def normalize_header(col):
    return re.sub(r"[^a-z0-9]+", "", col.strip().lower())


def normalize_species(value):
    v = re.sub(r"[^a-z]+", " ", value.strip().lower())
    v = re.sub(r"\s+", " ", v).strip()
    if v in {"human", "homo sapiens", "homo sapienst"}:
        return "homo sapiens"
    return v


def download_url_to_file(url, dest_path, allow_insecure=False):
    os.makedirs(os.path.dirname(dest_path) or ".", exist_ok=True)
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    context = ssl._create_unverified_context() if allow_insecure else None
    if context:
        response = urllib.request.urlopen(req, context=context)
    else:
        response = urllib.request.urlopen(req)
    with response as resp, open(dest_path, "wb") as f:
        while True:
            chunk = resp.read(1024 * 1024)
            if not chunk:
                break
            f.write(chunk)


def find_mirtarbase_mti_url(download_page_url):
    req = urllib.request.Request(download_page_url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req) as resp:
        html = resp.read().decode("utf-8", errors="ignore")
    links = re.findall(r'href=["\\\']([^"\\\']+)["\\\']', html, flags=re.IGNORECASE)
    scored = []
    for link in links:
        if "mirtarbase" not in link.lower():
            continue
        if "mti" not in link.lower():
            continue
        if not re.search(r"\\.(txt|tsv|csv|xls|xlsx|zip)$", link, flags=re.IGNORECASE):
            continue
        scored.append(link)
    if not scored:
        return None
    scored.sort(key=lambda s: (0 if s.lower().endswith(".txt") else 1, len(s)))
    return urllib.parse.urljoin(download_page_url, scored[0])


def read_mirtarbase(
    path,
    species="Homo sapiens",
    support=None,
    sheet=None,
    include_mirnas=None,
):
    species_norm = normalize_species(species) if species else None
    ext = os.path.splitext(path)[1].lower()
    if ext == ".xlsx":
        if sheet is None:
            with zipfile.ZipFile(path) as zf:
                sheets = list_sheets(zf)
                if not sheets:
                    raise ValueError("miRTarBase .xlsx has no sheets.")
                sheet = sheets[0][0]
        rows = read_xlsx_sheet(path, sheet)
        if not rows:
            return {"raw": defaultdict(set), "canonical": defaultdict(set)}
        header = rows[0]
        row_iter = iter(rows[1:])
    else:
        row_iter = iter_delimited_rows(path)
        try:
            header = next(row_iter)
        except StopIteration:
            return {"raw": defaultdict(set), "canonical": defaultdict(set)}

    idx_mirna = None
    idx_target = None
    idx_species_target = None
    idx_species_mirna = None
    idx_support = None

    for i, col in enumerate(header):
        key = normalize_header(col)
        if key in {"mirna", "mirnaname", "mirnaid", "mirname"} and idx_mirna is None:
            idx_mirna = i
        if key in {"targetgene", "target", "gene", "genesymbol"} and idx_target is None:
            idx_target = i
        if key in {"speciestargetgene", "speciestarget", "targetspecies"} and idx_species_target is None:
            idx_species_target = i
        if key in {"speciesmirna", "speciessourcemiRNA", "speciesmir"} and idx_species_mirna is None:
            idx_species_mirna = i
        if key in {"supporttype", "support", "evidence"} and idx_support is None:
            idx_support = i

    if idx_mirna is None or idx_target is None:
        raise ValueError(
            "miRTarBase format not recognized. Expected columns like 'miRNA' and 'Target Gene'."
        )

    mapping = defaultdict(set)
    support_norm = support.strip().lower() if support else None

    include_canonical = None
    if include_mirnas:
        include_canonical = {canonical_mirna(m) for m in include_mirnas}

    for row in row_iter:
        if len(row) <= max(idx_mirna, idx_target):
            continue
        mirna = row[idx_mirna].strip()
        target = row[idx_target].strip()
        if not mirna or not target:
            continue

        if include_canonical is not None and canonical_mirna(mirna) not in include_canonical:
            continue

        if species_norm:
            species_value = ""
            if idx_species_target is not None and idx_species_target < len(row):
                species_value = row[idx_species_target]
            elif idx_species_mirna is not None and idx_species_mirna < len(row):
                species_value = row[idx_species_mirna]
            if species_value:
                if normalize_species(species_value) != species_norm:
                    continue

        if support_norm and idx_support is not None and idx_support < len(row):
            if support_norm not in row[idx_support].lower():
                continue

        mapping[mirna].add(target)

    canonical = defaultdict(set)
    for mirna, targets in mapping.items():
        canonical[canonical_mirna(mirna)].update(targets)

    return {"raw": mapping, "canonical": canonical}
def parse_targets_rows(rows):
    if not rows:
        return {}

    header = rows[0]
    idx_mirna = None
    idx_target = None
    for i, col in enumerate(header):
        key = normalize_header(col)
        if key in {"mirna", "mir", "mirname", "mirid"} and idx_mirna is None:
            idx_mirna = i
        if key in {"target", "gene", "protein", "symbol", "targetgene", "targets"} and idx_target is None:
            idx_target = i

    if idx_mirna is None or idx_target is None:
        idx_mirna, idx_target = 0, 1

    mapping = defaultdict(set)
    for row in rows[1:]:
        if len(row) <= max(idx_mirna, idx_target):
            continue
        mirna = row[idx_mirna].strip()
        target_cell = row[idx_target].strip()
        if not mirna or not target_cell:
            continue
        targets = [target_cell]
        if ";" in target_cell:
            targets = [t.strip() for t in target_cell.split(";") if t.strip()]
        elif "," in target_cell:
            targets = [t.strip() for t in target_cell.split(",") if t.strip()]

        for target in targets:
            mapping[mirna].add(target)

    canonical = defaultdict(set)
    for mirna, targets in mapping.items():
        canonical[canonical_mirna(mirna)].update(targets)

    return {"raw": mapping, "canonical": canonical}


def select_high_expression(rows, metric, top_n, min_score, quantile):
    header = rows[0]
    sample_cols = header[1:]
    entries = []

    for row in rows[1:]:
        if not row or not row[0]:
            continue
        counts = []
        for val in row[1:]:
            try:
                counts.append(float(val))
            except Exception:
                counts.append(0.0)
        score = score_counts(counts, metric)
        entries.append((row[0], counts, score))

    entries.sort(key=lambda x: x[2], reverse=True)
    scores = [e[2] for e in entries]
    selected = entries

    if top_n is not None:
        selected = entries[: max(0, top_n)]
    elif quantile is not None:
        if not scores:
            selected = []
        else:
            idx = int(round((len(scores) - 1) * quantile))
            thresh = sorted(scores)[idx]
            selected = [e for e in entries if e[2] >= thresh]
    elif min_score is not None:
        selected = [e for e in entries if e[2] >= min_score]
    else:
        selected = entries[:50]

    return sample_cols, selected


def write_high_expression(path, sample_cols, selected, metric):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", newline="") as f:
        writer = csv.writer(f, delimiter="\t")
        writer.writerow(["miRNA", f"{metric}_score", *sample_cols])
        for mirna, counts, score in selected:
            writer.writerow([mirna, f"{score:.6g}", *counts])


def write_targets(outdir, selected, target_map, metric):
    edges_path = os.path.join(outdir, "miRNA_protein_edges.tsv")
    summary_path = os.path.join(outdir, "protein_targets_summary.tsv")
    os.makedirs(outdir, exist_ok=True)

    target_stats = defaultdict(lambda: {"miRNAs": set(), "score_sum": 0.0})
    with open(edges_path, "w", newline="") as f:
        writer = csv.writer(f, delimiter="\t")
        writer.writerow(["miRNA", "target", f"{metric}_score", "edge_weight"])
        for mirna, _counts, score in selected:
            targets = set()
            targets.update(target_map["raw"].get(mirna, set()))
            targets.update(target_map["canonical"].get(canonical_mirna(mirna), set()))
            for target in sorted(targets):
                writer.writerow([mirna, target, f"{score:.6g}", f"{score:.6g}"])
                target_stats[target]["miRNAs"].add(mirna)
                target_stats[target]["score_sum"] += score

    with open(summary_path, "w", newline="") as f:
        writer = csv.writer(f, delimiter="\t")
        writer.writerow(["target", "miRNA_count", f"{metric}_score_sum"])
        for target, stats in sorted(
            target_stats.items(), key=lambda item: (-len(item[1]["miRNAs"]), item[0])
        ):
            writer.writerow([target, len(stats["miRNAs"]), f"{stats['score_sum']:.6g}"])

    return edges_path, summary_path


def main():
    parser = argparse.ArgumentParser(description="Select high-expression miRNAs and build a target map.")
    parser.add_argument("--counts", default="miRNA_counts.xlsx", help="Path to miRNA counts .xlsx file.")
    parser.add_argument("--sheet", default="Mature", help="Sheet name to use (e.g., Mature or Hairpin).")
    parser.add_argument("--metric", choices=["mean", "median", "sum", "max"], default="mean")
    parser.add_argument("--top", type=int, default=None, help="Select top N miRNAs by metric.")
    parser.add_argument(
        "--min-score",
        type=float,
        default=5000.0,
        help="Select miRNAs with metric >= value. Default: 5000.",
    )
    parser.add_argument("--quantile", type=float, default=None, help="Select miRNAs with metric >= quantile (0-1).")
    parser.add_argument("--targets", default=None, help="TSV/CSV file with miRNA-to-target mappings.")
    parser.add_argument("--mirtarbase", default=None, help="miRTarBase TSV/CSV/XLSX file to build targets.")
    parser.add_argument(
        "--fetch-mirtarbase",
        action="store_true",
        help="Auto-download miRTarBase MTI file if missing (best-effort).",
    )
    parser.add_argument(
        "--mirtarbase-url",
        default=None,
        help="Direct URL to miRTarBase MTI file (overrides auto-detection).",
    )
    parser.add_argument(
        "--mirtarbase-download-page",
        default="http://mirtarbase.cuhk.edu.cn/php/download.php",
        help="miRTarBase download page URL to discover MTI file.",
    )
    parser.add_argument(
        "--mirtarbase-insecure",
        action="store_true",
        help="Allow insecure SSL (disable certificate verification) for downloads.",
    )
    parser.add_argument(
        "--mirtarbase-species",
        default="Homo sapiens",
        help="Species filter for miRTarBase (default: Homo sapiens).",
    )
    parser.add_argument(
        "--mirtarbase-support",
        default=None,
        help="Optional miRTarBase support filter (e.g., Strong).",
    )
    parser.add_argument(
        "--mirtarbase-sheet",
        default=None,
        help="Sheet name for miRTarBase .xlsx (defaults to first sheet).",
    )
    parser.add_argument("--outdir", default="out", help="Output directory.")
    args = parser.parse_args()

    if not os.path.exists(args.counts):
        print(f"Counts file not found: {args.counts}", file=sys.stderr)
        return 2

    try:
        rows = read_xlsx_sheet(args.counts, args.sheet)
    except Exception as exc:
        print(f"Failed to read sheet: {exc}", file=sys.stderr)
        return 2

    if not rows:
        print("No data found in sheet.", file=sys.stderr)
        return 2

    sample_cols, selected = select_high_expression(
        rows, args.metric, args.top, args.min_score, args.quantile
    )
    if not selected:
        print("No miRNAs selected with the current criteria.", file=sys.stderr)
        return 2

    outdir = args.outdir
    high_path = os.path.join(outdir, "high_expression_miRNAs.tsv")
    write_high_expression(high_path, sample_cols, selected, args.metric)

    edges_path = summary_path = None
    target_map = None
    if args.mirtarbase or args.fetch_mirtarbase:
        if args.mirtarbase is None:
            args.mirtarbase = os.path.join(args.outdir, "miRTarBase_MTI.txt")

        if not os.path.exists(args.mirtarbase) and args.fetch_mirtarbase:
            try:
                url = args.mirtarbase_url
                if url is None:
                    url = find_mirtarbase_mti_url(args.mirtarbase_download_page)
                if url is None:
                    print(
                        "Could not auto-detect miRTarBase MTI file URL from the download page.",
                        file=sys.stderr,
                    )
                    return 2
                print(f"Downloading miRTarBase MTI file from: {url}")
                try:
                    download_url_to_file(url, args.mirtarbase, allow_insecure=args.mirtarbase_insecure)
                except urllib.error.URLError as exc:
                    reason = getattr(exc, "reason", None)
                    if isinstance(reason, ssl.SSLError) and url.startswith("https://"):
                        http_url = "http://" + url[len("https://") :]
                        print(f"SSL failed, retrying over HTTP: {http_url}")
                        download_url_to_file(
                            http_url, args.mirtarbase, allow_insecure=args.mirtarbase_insecure
                        )
                    else:
                        raise
            except (urllib.error.URLError, ValueError, OSError) as exc:
                print(f"Failed to download miRTarBase file: {exc}", file=sys.stderr)
                return 2

        if not os.path.exists(args.mirtarbase):
            print(f"miRTarBase file not found: {args.mirtarbase}", file=sys.stderr)
            return 2

        try:
            target_map = read_mirtarbase(
                args.mirtarbase,
                species=args.mirtarbase_species,
                support=args.mirtarbase_support,
                sheet=args.mirtarbase_sheet,
                include_mirnas=[m for m, _c, _s in selected],
            )
        except Exception as exc:
            print(f"Failed to read miRTarBase file: {exc}", file=sys.stderr)
            return 2
    elif args.targets:
        if not os.path.exists(args.targets):
            print(f"Targets file not found: {args.targets}", file=sys.stderr)
            return 2
        target_map = read_targets(args.targets)

    if target_map:
        edges_path, summary_path = write_targets(outdir, selected, target_map, args.metric)

    print(f"Wrote high-expression list: {high_path}")
    if edges_path:
        print(f"Wrote target edges: {edges_path}")
        print(f"Wrote target summary: {summary_path}")
    else:
        print("No target map generated (provide --mirtarbase or --targets).")
    return 0


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