"""
FSC (Fragment Size Coverage) processor.

Processes GC-weighted fragment counts into ML-ready features per the implementation plan:
- Raw GC-weighted counts per channel (from Rust backend)
- PoN log2 ratios: log2(channel / PoN_mean)
- Reliability scores: 1 / (PoN_variance + k)

NO z-scores - the plan specifies GC-weighting (Rust) + PoN log-ratios (Python).
"""

from pathlib import Path
import pandas as pd
import numpy as np
import logging

logger = logging.getLogger("core.fsc_processor")

# 6 biologically-meaningful channels matching Rust backend
# Extended to include ultra_long for fragments 401-1000bp
CHANNELS = ["ultra_short", "core_short", "mono_nucl", "di_nucl", "long", "ultra_long"]

# Regularization constant for reliability calculation
RELIABILITY_K = 0.01


# ============================================================================
# GC CORRECTION FACTOR LOADING
# ============================================================================
#
# For gene-level FSC in panel mode, we apply GC correction weights to account
# for capture bias. Different GC content leads to different hybridization
# efficiency and PCR amplification rates.
#
# The correction factors are loaded from .correction_factors.ontarget.tsv
# generated by the extract step with --target-regions.
# ============================================================================


def load_correction_factors(factor_path: Path) -> dict:
    """
    Load GC correction factors into a fast lookup table.

    The factors are indexed by (length_bin, gc_percent) for O(1) lookup.
    Length bins match the Rust backend: bin = (fragment_length - 60) // 5

    Args:
        factor_path: Path to correction_factors.ontarget.tsv or .tsv

    Returns:
        Dict[(len_bin, gc_pct), factor] for fast lookup
        Empty dict if loading fails (falls back to unweighted counting)
    """
    factors = {}

    try:
        # Sniff actual delimiter from first line (files may have wrong extension)
        with open(factor_path, "r") as f:
            first_line = f.readline()
        sep = "," if "," in first_line else "\t"
        df = pd.read_csv(factor_path, sep=sep)

        # Expected columns: length_bin (or similar), gc_percent, factor
        # Find the relevant columns
        len_col = None
        gc_col = None
        factor_col = None

        for col in df.columns:
            col_lower = col.lower()
            if "length" in col_lower and "bin" in col_lower:
                len_col = col
            elif "gc" in col_lower and ("percent" in col_lower or "pct" in col_lower):
                gc_col = col
            elif "factor" in col_lower or "correction" in col_lower:
                factor_col = col

        if not all([len_col, gc_col, factor_col]):
            logger.warning(
                f"Correction factors file missing expected columns: {list(df.columns)}"
            )
            return factors

        # Build lookup table
        # Note: length_bin column may contain length_bin_min values (e.g., 75, 80, ...)
        # or bin indices (e.g., 3, 4, ...). We detect which format and convert.
        is_raw_lengths = len_col.lower().endswith("_min") or df[len_col].min() >= 60

        for _, row in df.iterrows():
            try:
                len_val = int(row[len_col])
                gc_pct = int(row[gc_col])
                factor = float(row[factor_col])

                # Convert raw length to bin index if needed
                # Bin formula matches Rust: (length - 60) // 5
                if is_raw_lengths:
                    len_bin = (len_val - 60) // 5
                else:
                    len_bin = len_val

                # Validate bin index (0-67 for lengths 60-399)
                if len_bin < 0 or len_bin > 67:
                    continue

                # Validate factor (should be positive, not too extreme)
                if 0.01 < factor < 100:
                    factors[(len_bin, gc_pct)] = factor
            except (ValueError, TypeError):
                continue

        logger.info(
            f"Loaded {len(factors)} GC correction factors from {factor_path.name}"
        )

    except FileNotFoundError:
        logger.warning(f"Correction factors file not found: {factor_path}")
    except Exception as e:
        logger.warning(f"Failed to load correction factors from {factor_path}: {e}")

    return factors


def _get_gc_weight(frag_len: int, gc_frac: float, correction_factors: dict) -> float:
    """
    Get the GC correction weight for a fragment.

    Args:
        frag_len: Fragment length in bp
        gc_frac: GC fraction (0.0 to 1.0)
        correction_factors: Lookup table from load_correction_factors()

    Returns:
        Correction weight (default 1.0 if not found)
    """
    if not correction_factors:
        return 1.0

    # Calculate length bin (matches Rust: (len - 60) / 5, clamped to 0-67)
    safe_len = max(60, min(399, frag_len))
    len_bin = (safe_len - 60) // 5

    # Convert GC fraction to percentage (0-100)
    gc_pct = int(round(gc_frac * 100))
    gc_pct = max(0, min(100, gc_pct))

    return correction_factors.get((len_bin, gc_pct), 1.0)


def process_fsc(
    counts_df: pd.DataFrame,
    output_path: Path,
    windows: int = 100000,
    continue_n: int = 50,
    pon=None,
) -> Path:
    """
    Process GC-weighted fragment counts into FSC ML features.

    Per the implementation plan, FSC normalization order is:
    1. GC-weight (Rust) - done upstream
    2. PoN log-ratio (Python) - done here

    Output columns:
    - chrom, start, end: Bin coordinates
    - ultra_short, core_short, mono_nucl, di_nucl, long: GC-weighted counts
    - total: All fragments
    - *_log2: log2(channel / PoN_mean) when PoN provided
    - *_reliability: 1 / (PoN_variance + k) when PoN provided

    Args:
        counts_df: DataFrame from Rust with [chrom, start, end, ultra_short,
                   core_short, mono_nucl, di_nucl, long, total, mean_gc]
        output_path: Path to write FSC.tsv output
        windows: Window size in bp (default 100000)
        continue_n: Number of bins to aggregate (default 50)
        pon: Optional PON model for log2 ratio normalization

    Returns:
        Path to the output file
    """
    logger.info(f"Processing FSC: {len(counts_df)} bins → {output_path}")

    # Validate required columns
    required_cols = ["chrom", "start", "end"] + CHANNELS + ["total"]
    missing = [c for c in required_cols if c not in counts_df.columns]
    if missing:
        raise ValueError(f"Missing required columns: {missing}")

    # Aggregate bins into windows
    results = []
    for chrom, group in counts_df.groupby("chrom", sort=False):
        n_bins = len(group)
        n_windows = n_bins // continue_n

        if n_windows == 0:
            continue

        trunc_len = n_windows * continue_n

        # Aggregate each channel
        channel_sums = {}
        for ch in CHANNELS:
            mat = group[ch].values[:trunc_len].reshape(n_windows, continue_n)
            channel_sums[ch] = mat.sum(axis=1)

        totals_mat = group["total"].values[:trunc_len].reshape(n_windows, continue_n)
        channel_sums["total"] = totals_mat.sum(axis=1)

        # Get window coordinates from first and last bin in each window
        starts = group["start"].values[:trunc_len].reshape(n_windows, continue_n)[:, 0]
        ends = group["end"].values[:trunc_len].reshape(n_windows, continue_n)[:, -1]

        window_df = pd.DataFrame(
            {"chrom": chrom, "start": starts, "end": ends, **channel_sums}
        )
        results.append(window_df)

    if not results:
        logger.warning("No valid windows found for FSC")
        _write_empty_fsc(output_path, pon is not None)
        return output_path

    final_df = pd.concat(results, ignore_index=True)
    logger.info(f"Aggregated into {len(final_df)} windows")

    # Compute PoN log2 ratios and reliability if model provided
    if pon is not None:
        logger.info("Computing PoN log2 ratios...")
        for ch in CHANNELS:
            mean = pon.get_mean(ch) if hasattr(pon, "get_mean") else None
            var = pon.get_variance(ch) if hasattr(pon, "get_variance") else None

            if mean and mean > 0:
                # log2(channel / PoN_mean)
                final_df[f"{ch}_log2"] = np.log2(final_df[ch] / mean + 1e-9)
            else:
                final_df[f"{ch}_log2"] = 0.0

            if var is not None:
                # reliability = 1 / (PoN_variance + k)
                final_df[f"{ch}_reliability"] = 1.0 / (var + RELIABILITY_K)
            else:
                final_df[f"{ch}_reliability"] = 1.0

    # Write output
    _write_fsc_output(final_df, output_path, pon is not None)

    logger.info(f"FSC complete: {output_path}")
    return output_path


def _write_empty_fsc(output_path: Path, with_pon: bool):
    """Write empty FSC file with appropriate headers."""
    cols = ["chrom", "start", "end"] + CHANNELS + ["total"]
    if with_pon:
        cols += [f"{ch}_log2" for ch in CHANNELS]
        cols += [f"{ch}_reliability" for ch in CHANNELS]

    pd.DataFrame(columns=cols).to_csv(output_path, sep="\t", index=False)


def _write_fsc_output(df: pd.DataFrame, output_path: Path, with_pon: bool):
    """Write FSC output TSV matching the implementation plan schema."""
    # Determine columns to output
    cols = ["chrom", "start", "end"] + CHANNELS + ["total"]
    if with_pon:
        cols += [f"{ch}_log2" for ch in CHANNELS]
        cols += [f"{ch}_reliability" for ch in CHANNELS]

    # Select and write
    output_df = df[cols].copy()
    output_df.to_csv(output_path, sep="\t", index=False, float_format="%.4f")


def aggregate_by_gene(
    fragments_bed: Path,
    genes: dict,
    output_path: Path,
    pon=None,
    correction_factors: dict = None,
    aggregate_by: str = "gene",
    gene_bed_path: Path = None,
) -> Path:
    """
    Aggregate fragment counts for panel-specific FSC using Rust implementation.

    This function counts fragments per gene using the high-performance Rust engine,
    applying GC correction if factors are available.

    Args:
        fragments_bed: Path to fragment BED file from extract step
        genes: Dict[gene_name, List[Region]] from gene_bed.load_gene_bed() (used only if gene_bed_path not provided)
        output_path: Path to write FSC output
        pon: Optional PON model for normalization (legacy, not used in Rust)
        correction_factors: Optional dict or Path to GC correction factors
        aggregate_by: 'gene' (sum by gene) or 'region' (one row per exon/region)
        gene_bed_path: Path to gene BED file (if provided, used directly by Rust)

    Returns:
        Path to output file

    Output columns (gene mode):
        - gene, n_regions, total_bp, channels, *_ratio, normalized_depth
    """
    from krewlyzer import _core
    import tempfile

    logger.info(f"Aggregating FSC by {aggregate_by}: {len(genes)} genes")

    # Prepare gene BED for Rust
    if gene_bed_path and Path(gene_bed_path).exists():
        # Use provided gene BED directly
        use_gene_bed = Path(gene_bed_path)
    else:
        # Write genes dict to temp BED for Rust
        import tempfile

        with tempfile.NamedTemporaryFile(mode="w", suffix=".bed", delete=False) as f:
            for gene, regions in genes.items():
                for region in regions:
                    name = getattr(region, "name", gene)
                    f.write(
                        f"{region.chrom}\t{region.start}\t{region.end}\t{gene}\t{name}\n"
                    )
            use_gene_bed = Path(f.name)

    # Prepare GC factors path
    gc_factors_path = None
    if correction_factors:
        if isinstance(correction_factors, (str, Path)):
            gc_factors_path = str(correction_factors)
        elif isinstance(correction_factors, dict):
            # Write factors dict to temp TSV for Rust
            with tempfile.NamedTemporaryFile(
                mode="w", suffix=".tsv", delete=False
            ) as f:
                f.write("length_bin\tgc_pct\tfactor\n")
                for (len_bin, gc_pct), factor in correction_factors.items():
                    f.write(f"{len_bin}\t{gc_pct}\t{factor}\n")
                gc_factors_path = f.name

    # Call Rust implementation
    n_output = _core.aggregate_by_gene(
        str(fragments_bed),
        str(use_gene_bed),
        str(output_path),
        gc_factors_path,
        aggregate_by,
    )

    # Cleanup temp file if we created one
    if not gene_bed_path or not Path(gene_bed_path).exists():
        try:
            use_gene_bed.unlink(missing_ok=True)
        except Exception:
            pass

    logger.info(f"FSC {aggregate_by} complete: {output_path} ({n_output} rows)")
    return output_path


def apply_fsc_gene_pon(
    fsc_gene_path: Path,
    pon,
    output_path: Path = None,
) -> Path:
    """
    Apply FSC gene PON baseline z-score normalization.

    Reads FSC.gene.tsv, computes z-score for each gene's normalized_depth
    against the PON baseline, and adds a depth_zscore column.

    Args:
        fsc_gene_path: Path to FSC.gene.tsv file
        pon: PonModel with fsc_gene_baseline
        output_path: Output path (defaults to overwriting input)

    Returns:
        Path to output file with depth_zscore column

    Note:
        Clinical interpretation:
        - z-score >> 0: Gene amplification
        - z-score << 0: Gene deletion
        - z-score ≈ 0: Normal copy number
    """
    import pandas as pd

    if output_path is None:
        output_path = fsc_gene_path

    if not Path(fsc_gene_path).exists():
        logger.warning(f"FSC gene file not found: {fsc_gene_path}")
        return None

    if not pon or not pon.fsc_gene_baseline:
        logger.debug("No FSC gene baseline in PON, skipping z-score")
        return fsc_gene_path

    df = pd.read_csv(fsc_gene_path, sep="\t")

    if "gene" not in df.columns or "normalized_depth" not in df.columns:
        logger.warning(f"FSC gene file missing required columns: {fsc_gene_path}")
        return fsc_gene_path

    # Compute z-scores using PON baseline
    baseline = pon.fsc_gene_baseline
    z_scores = []
    for _, row in df.iterrows():
        gene = row["gene"]
        depth = row["normalized_depth"]
        z = baseline.compute_zscore(gene, depth)
        z_scores.append(z if z is not None else float("nan"))

    df["depth_zscore"] = z_scores

    # Write output
    df.to_csv(output_path, sep="\t", index=False)
    n_with_zscore = df["depth_zscore"].notna().sum()
    logger.debug(f"Applied FSC gene PON: {n_with_zscore}/{len(df)} genes with z-scores")

    return output_path


def apply_fsc_region_pon(
    fsc_region_path: Path,
    pon,
    output_path: Path = None,
) -> Path:
    """
    Apply FSC region (exon/probe) PON baseline z-score normalization.

    Reads FSC.regions.tsv, computes z-score for each region's normalized_depth
    against the PON baseline, and adds a depth_zscore column.

    Args:
        fsc_region_path: Path to FSC.regions.tsv file
        pon: PonModel with fsc_region_baseline
        output_path: Output path (defaults to overwriting input)

    Returns:
        Path to output file with depth_zscore column

    Note:
        Region IDs are formatted as "chrom:start-end".
        Clinical interpretation:
        - z-score << 0: Exon deletion
        - z-score >> 0: Exon amplification
    """
    import pandas as pd

    if output_path is None:
        output_path = fsc_region_path

    if not Path(fsc_region_path).exists():
        logger.warning(f"FSC region file not found: {fsc_region_path}")
        return None

    if not pon or not pon.fsc_region_baseline:
        logger.debug("No FSC region baseline in PON, skipping z-score")
        return fsc_region_path

    df = pd.read_csv(fsc_region_path, sep="\t")

    required_cols = ["chrom", "start", "end", "normalized_depth"]
    if not all(c in df.columns for c in required_cols):
        logger.warning(f"FSC region file missing required columns: {fsc_region_path}")
        return fsc_region_path

    # Create region_id for lookup
    df["region_id"] = (
        df["chrom"].astype(str)
        + ":"
        + df["start"].astype(str)
        + "-"
        + df["end"].astype(str)
    )

    # Compute z-scores using PON baseline
    baseline = pon.fsc_region_baseline
    z_scores = []
    for _, row in df.iterrows():
        region_id = row["region_id"]
        depth = row["normalized_depth"]
        z = baseline.compute_zscore(region_id, depth)
        z_scores.append(z if z is not None else float("nan"))

    df["depth_zscore"] = z_scores

    # Drop temporary column
    df = df.drop(columns=["region_id"])

    # Write output
    df.to_csv(output_path, sep="\t", index=False)
    n_with_zscore = df["depth_zscore"].notna().sum()
    logger.debug(
        f"Applied FSC region PON: {n_with_zscore}/{len(df)} regions with z-scores"
    )

    return output_path


# ============================================================================
# E1-ONLY FILTERING
# ============================================================================
#
# First exon (E1) filtering extracts only promoter-proximal regions.
# Per Helzer et al. (2025), promoters are Nucleosome Depleted Regions (NDRs)
# with distinct fragmentation patterns. E1 signal is often stronger for
# early cancer detection than whole-gene averages.
#
# E1 is identified as the first exon by genomic position per gene,
# consistent with the region-mds implementation.
# ============================================================================


def filter_fsc_to_e1(
    fsc_regions_path: Path,
    output_path: Path = None,
) -> Path:
    """
    Filter FSC regions to E1 (first exon) per gene.

    Extracts the first exon by genomic position for each gene from the
    full FSC.regions.tsv file. This focuses on promoter-proximal regions
    where cancer fragmentation signal is strongest.

    Args:
        fsc_regions_path: Path to FSC.regions.tsv file
        output_path: Output path (defaults to input.e1only.tsv)

    Returns:
        Path to E1-only output file, or None if input doesn't exist

    Scientific context:
        E1 (first exon) serves as a proxy for promoter regions which are
        Nucleosome Depleted Regions (NDRs). Differential enzymatic access
        at NDRs produces distinct fragment end motifs. Cancer signal is
        often stronger at E1 than averaging across all exons.
    """
    import pandas as pd

    fsc_regions_path = Path(fsc_regions_path)

    if not fsc_regions_path.exists():
        logger.warning(f"FSC regions file not found: {fsc_regions_path}")
        return None

    # Default output path: replace .tsv with .e1only.tsv
    if output_path is None:
        stem = fsc_regions_path.stem
        if stem.endswith(".regions"):
            stem = stem.replace(".regions", ".regions.e1only")
        else:
            stem = f"{stem}.e1only"
        output_path = fsc_regions_path.parent / f"{stem}.tsv"
    output_path = Path(output_path)

    # Read input
    df = pd.read_csv(fsc_regions_path, sep="\t")

    # Validate required columns
    required_cols = ["gene", "chrom", "start"]
    missing = [c for c in required_cols if c not in df.columns]
    if missing:
        logger.warning(
            f"FSC regions file missing required columns {missing}: {fsc_regions_path}"
        )
        return None

    if df.empty:
        logger.warning(f"FSC regions file is empty: {fsc_regions_path}")
        return None

    n_input = len(df)
    n_genes = df["gene"].nunique()

    # Sort by gene + genomic position, take first per gene (E1 = lowest start)
    df_sorted = df.sort_values(["gene", "chrom", "start"])
    e1_df = df_sorted.groupby("gene", as_index=False).first()

    # Preserve original column order
    original_cols = [c for c in df.columns if c in e1_df.columns]
    e1_df = e1_df[original_cols]

    # Write output
    e1_df.to_csv(output_path, sep="\t", index=False)

    logger.info(
        f"FSC E1-only filter: {n_input} regions → {len(e1_df)} E1 regions ({n_genes} genes) → {output_path.name}"
    )

    return output_path
