#!/usr/bin/env python3
"""Post-discovery Project 3087 channel / bin-phase / flight-height sensitivity.

Usage
-----
python verify_3087_channel_altitude_sensitivity.py --input /path/to/Mag-3087.XYZ
python verify_3087_channel_altitude_sensitivity.py --input /path/to/3087-base.zip

The script streams the public upstream Project 3087 flight-line file, keeps only
14 frozen Uruana production lines and the pre-existing target/control windows,
and recomputes a deliberately separate sensitivity family.  It does NOT alter
the manuscript's frozen primary MAGBRU operator or its reported line slopes.

No upstream raw observations are redistributed in Online Resource 2.
"""
from __future__ import annotations
import argparse, io, zipfile
from pathlib import Path
import numpy as np
import pandas as pd

CX = 365408.554
CY = 8224103.341
LINES = set(range(12340, 12471, 10))
MAG_CHANNELS = ["MAGBRU", "MAGCOM", "MAGCOR", "MAGNIV", "MAGMIC", "IGRF"]
HEIGHT_CHANNELS = ["GPSALT", "BARO", "ALTURA", "MDT"]
KEEP = ["X", "Y"] + MAG_CHANNELS + HEIGHT_CHANNELS
PHASES_M = [0.0, 62.5, 125.0, 187.5]


def _open_text(path: Path):
    if zipfile.is_zipfile(path):
        z = zipfile.ZipFile(path)
        matches = [n for n in z.namelist() if n.replace('\\','/').endswith('/XYZ/Mag-3087.XYZ') or n.endswith('Mag-3087.XYZ')]
        if len(matches) != 1:
            z.close()
            raise RuntimeError(f"Expected one Mag-3087.XYZ member, found {len(matches)}")
        raw = z.open(matches[0], 'r')
        txt = io.TextIOWrapper(raw, encoding='latin1', errors='replace', newline='')
        return txt, z
    return open(path, 'r', encoding='latin1', errors='replace', newline=''), None


def load_relevant(path: Path) -> pd.DataFrame:
    txt, owner = _open_text(path)
    columns = None
    idx = None
    current_line = None
    rows = []
    ylo, yhi = CY + 6000.0, CY + 20000.0
    try:
        for raw in txt:
            s = raw.strip()
            if not s:
                continue
            if s.startswith('/X,'):
                columns = s[1:].split(',')
                idx = {c:i for i,c in enumerate(columns)}
                missing = [c for c in KEEP if c not in idx]
                if missing:
                    raise RuntimeError(f"Missing expected columns: {missing}")
                continue
            if s.lower().startswith('line'):
                try:
                    current_line = int(s.split()[-1])
                except Exception:
                    current_line = None
                continue
            if s.startswith('/') or s.startswith('//') or current_line not in LINES or idx is None:
                continue
            parts = s.split(',')
            try:
                y = float(parts[idx['Y']])
            except Exception:
                continue
            if not (ylo <= y <= yhi):
                continue
            rec = {'line': current_line}
            for c in KEEP:
                v = parts[idx[c]] if idx[c] < len(parts) else '*'
                try: rec[c] = float(v) if v not in ('*','') else np.nan
                except Exception: rec[c] = np.nan
            rows.append(rec)
    finally:
        txt.close()
        if owner is not None: owner.close()
    d = pd.DataFrame(rows)
    if d.empty:
        raise RuntimeError('No relevant rows recovered from upstream file.')
    d['dy_m'] = d['Y'] - CY
    d['r_m'] = np.hypot(d['X'] - CX, d['Y'] - CY)
    got = set(d['line'].astype(int).unique())
    if got != LINES:
        raise RuntimeError(f"Expected frozen lines {sorted(LINES)}, recovered {sorted(got)}")
    return d


def pairwise_median_slope(x, y):
    x = np.asarray(x, float); y = np.asarray(y, float)
    slopes = []
    for i in range(len(x)-1):
        dx = x[i+1:] - x[i]
        ok = dx != 0
        slopes.extend(((y[i+1:][ok] - y[i]) / dx[ok]).tolist())
    return float(np.median(slopes)) if slopes else np.nan


def fit_slope(b: pd.DataFrame, col: str, estimator='ols') -> float:
    z = b[['Y',col]].dropna()
    if len(z) < 3: return np.nan
    x = (z['Y'].to_numpy(float) - CY) / 1000.0
    y = z[col].to_numpy(float)
    if estimator == 'ols': return float(np.polyfit(x, y, 1)[0])
    if estimator == 'theil': return pairwise_median_slope(x, y)
    raise ValueError(estimator)


def binned(g: pd.DataFrame, region: str, phase_m: float) -> pd.DataFrame:
    if region == 'target':
        q = g[(g.dy_m >= 6000.0) & (g.dy_m <= 8000.0) & (g.r_m <= 8000.0)].copy()
    else:
        q = g[(g.dy_m >= 16000.0) & (g.dy_m <= 20000.0)].copy()
    q['bin_id'] = np.floor((q['Y'] - phase_m) / 250.0).astype('int64')
    return q.groupby('bin_id')[['Y'] + MAG_CHANNELS + HEIGHT_CHANNELS].median(numeric_only=True).dropna(subset=['Y'])


def compute(d: pd.DataFrame):
    rows = []
    for phase in PHASES_M:
        for line, g in d.groupby('line'):
            for region in ('target','north'):
                b = binned(g, region, phase)
                row = {'phase_m':phase, 'line':int(line), 'region':region, 'n_bins':len(b)}
                for c in MAG_CHANNELS:
                    row[f'{c}_ols'] = fit_slope(b,c,'ols')
                    row[f'{c}_theil'] = fit_slope(b,c,'theil')
                for c in HEIGHT_CHANNELS:
                    row[f'{c}_ols'] = fit_slope(b,c,'ols')
                rows.append(row)
    line_table = pd.DataFrame(rows)

    summaries = []
    for phase in PHASES_M:
        for region in ('target','north'):
            z=line_table[(line_table.phase_m==phase)&(line_table.region==region)]
            for c in MAG_CHANNELS:
                for est in ('ols','theil'):
                    v=z[f'{c}_{est}']
                    summaries.append({'phase_m':phase,'region':region,'channel':c,'estimator':est,
                        'n_negative':int((v<0).sum()),'n_positive':int((v>0).sum()),'n':int(v.notna().sum()),
                        'median_slope':float(v.median())})
            for c in HEIGHT_CHANNELS:
                v=z[f'{c}_ols']
                summaries.append({'phase_m':phase,'region':region,'channel':c,'estimator':'ols',
                    'n_negative':int((v<0).sum()),'n_positive':int((v>0).sum()),'n':int(v.notna().sum()),
                    'median_slope':float(v.median())})
    return line_table, pd.DataFrame(summaries)


def rawpoint_summary(d: pd.DataFrame):
    rows=[]
    for line,g in d.groupby('line'):
        for region in ('target','north'):
            q = g[(g.dy_m>=6000)&(g.dy_m<=8000)&(g.r_m<=8000)] if region=='target' else g[(g.dy_m>=16000)&(g.dy_m<=20000)]
            row={'line':int(line),'region':region,'n_points':len(q)}
            for c in MAG_CHANNELS + HEIGHT_CHANNELS:
                z=q[['Y',c]].dropna(); x=(z.Y.to_numpy(float)-CY)/1000.; y=z[c].to_numpy(float)
                row[f'{c}_ols']=float(np.polyfit(x,y,1)[0]) if len(z)>=3 else np.nan
            rows.append(row)
    return pd.DataFrame(rows)


def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--input', required=True, help='Mag-3087.XYZ or ZIP archive containing it')
    ap.add_argument('--outdir', default=None, help='Optional directory to write recomputed CSVs')
    a=ap.parse_args()
    d=load_relevant(Path(a.input))
    line_table, summary=compute(d)
    raw=rawpoint_summary(d)
    if a.outdir:
        out=Path(a.outdir); out.mkdir(parents=True,exist_ok=True)
        line_table.to_csv(out/'URUANA_2011_BIN_PHASE_CHANNEL_HEIGHT_SENSITIVITY.csv',index=False)
        summary.to_csv(out/'URUANA_2011_BIN_PHASE_CHANNEL_HEIGHT_SENSITIVITY_SUMMARY.csv',index=False)
        raw.to_csv(out/'URUANA_2011_RAWPOINT_CHANNEL_HEIGHT_SENSITIVITY.csv',index=False)

    # Frozen post-discovery sensitivity expectations.
    for c in ('MAGCOM','MAGNIV','MAGMIC'):
        z=summary[(summary.channel==c)&(summary.estimator=='ols')]
        assert ((z[z.region=='target'].n_negative==14)).all()
        assert ((z[z.region=='north'].n_positive==14)).all()
    # MAGBRU and MAGCOR each expose one 13/14 target count at the 125-m phase; all north controls stay 14/14 positive.
    for c in ('MAGBRU','MAGCOR'):
        z=summary[(summary.channel==c)&(summary.estimator=='ols')]
        assert z[z.region=='target'].n_negative.min()==13
        assert z[z.region=='target'].n_negative.max()==14
        assert (z[z.region=='north'].n_positive==14).all()
    # Flight-height/topography channels do not duplicate a universal negative magnetic sign across target lines.
    for c,npos in [('GPSALT',13),('BARO',13),('ALTURA',4),('MDT',10)]:
        z=summary[(summary.region=='target')&(summary.channel==c)&(summary.estimator=='ols')]
        assert (z.n_positive==npos).all()
    # Without binning, all supplied magnetic-stage channels retain the target-negative / north-positive turnover.
    for c in ('MAGBRU','MAGCOM','MAGCOR','MAGNIV','MAGMIC'):
        assert (raw[raw.region=='target'][c+'_ols']<0).all()
        assert (raw[raw.region=='north'][c+'_ols']>0).all()

    print('Project 3087 post-discovery channel/height sensitivity: PASS')
    print('Rows retained from upstream XYZ:', len(d))
    print('MAGCOM/MAGNIV/MAGMIC: 14/14 target-negative and 14/14 north-positive at all four 250-m bin phases (OLS).')
    print('MAGBRU/MAGCOR: 13/14 at one 125-m shifted target phase, 14/14 at the other three; north remains 14/14 positive.')
    print('Target height/topography positive-slope counts across every phase: GPSALT 13/14; BARO 13/14; ALTURA 4/14; MDT 10/14.')
    print('Raw-point OLS: MAGBRU/MAGCOM/MAGCOR/MAGNIV/MAGMIC all 14/14 target-negative and 14/14 north-positive.')
    print('Interpretation: supports multi-channel turnover and rejects a simple universal sign mimic by the supplied height/topography channels; it is not an altitude-corrected replacement for the frozen primary operator.')

if __name__=='__main__': main()
