"""
Reimplementation of the 'Significant Window' algorithm in Davidson, U-Yen & Holloman (2015),
New Concepts in Global Tectonics 3(3), 310-322, Appendix A.

Inputs: WSO 10-day polar field values (unfiltered N and S, microtesla) and a USGS catalogue.
Ambiguous rules are exposed as options so they can be calibrated against the paper's own
reported in-sample statistics before any out-of-sample quake is looked at.

"""
import re, csv, datetime as dt
import numpy as np

MAX, MIN, REC = 0, 1, 2
STEPS_15YR = 55  # 1.5 years in 10-day steps


# ----------------------------------------------------------------------------- data
def load_wso(path, end=None):
    pat = re.compile(r'^(\d{4}):(\d{2}):(\d{2})_\S+\s+(-?\d+|XXX)N\s+(-?\d+|XXX)S\s+(-?\d+|XXX)Avg')
    dates, N, S = [], [], []
    for line in open(path):
        m = pat.match(line.strip())
        if not m:
            continue
        d = dt.date(int(m.group(1)), int(m.group(2)), int(m.group(3)))
        if end and d > end:
            break
        n, s = m.group(4), m.group(5)
        dates.append(d)
        N.append(np.nan if n == 'XXX' else float(n))
        S.append(np.nan if s == 'XXX' else float(s))
    dates, N, S = np.array(dates), np.array(N), np.array(S)
    # drop trailing missing values, interpolate interior gaps
    last = max(np.where(~np.isnan(N))[0].max(), np.where(~np.isnan(S))[0].max())
    dates, N, S = dates[:last + 1], N[:last + 1], S[:last + 1]
    n_interp = int(np.isnan(N).sum() + np.isnan(S).sum())
    idx = np.arange(len(N))
    for arr in (N, S):
        bad = np.isnan(arr)
        if bad.any():
            arr[bad] = np.round(np.interp(idx[bad], idx[~bad], arr[~bad]))
    return dates, N, S, n_interp


def load_quakes(path):
    out = []
    for r in csv.DictReader(open(path)):
        out.append((dt.date.fromisoformat(r['time'][:10]), float(r['mag']), r['place']))
    return out


# ----------------------------------------------------------------------------- smoothing
def longterm_extremes(X):
    """A_r = mean(X[r-100 : r+99]); extremes = points equal to the max/min of A within +-100."""
    R = len(X)
    A = np.array([X[max(0, r - 100):min(R, r + 100)].mean() for r in range(R)])
    ext = []
    for r in range(R):
        w = A[max(0, r - 100):min(R, r + 100)]
        if A[r] == w.max():
            ext.append((r, +1))
        elif A[r] == w.min():
            ext.append((r, -1))
    return A, ext


def minicycle_smooth(X):
    """Local trend + 1-yr sinusoid fits over +-100 steps; U_r = mean of all local fits at r."""
    R = len(X)
    acc, cnt = np.zeros(R), np.zeros(R)
    for c in range(R):
        lo, hi = max(0, c - 100), min(R, c + 100)
        i = np.arange(lo, hi)
        D = np.column_stack([np.ones_like(i, float), i, np.sin(2 * np.pi * 10 * i / 365.0),
                             np.cos(2 * np.pi * 10 * i / 365.0)])
        beta, *_ = np.linalg.lstsq(D, X[lo:hi], rcond=None)
        acc[lo:hi] += D @ beta
        cnt[lo:hi] += 1
    return acc / cnt


def local_extrema(U):
    ex = []
    for r in range(1, len(U) - 1):
        if U[r] > U[r - 1] and U[r] > U[r + 1]:
            ex.append((r, +1))
        elif U[r] < U[r - 1] and U[r] < U[r + 1]:
            ex.append((r, -1))
    return ex


# ----------------------------------------------------------------------------- half-minicycles and spikes
def half_minicycles(X, U):
    ex = local_extrema(U)
    segs = []
    for (r1, t1), (r2, t2) in zip(ex[:-1], ex[1:]):
        target = 0.5 * (X[r1] + X[r2])
        rng = np.arange(r1, r2 + 1)
        segs.append(int(rng[np.argmin(np.abs(X[rng] - target))]))
    hcs = []
    for k in range(len(segs) - 1):
        a, b = segs[k], segs[k + 1]
        inside = [t for (r, t) in ex if a <= r <= b]
        if not inside:
            continue
        hcs.append({'start': a, 'end': b, 's': inside[0]})
    return hcs


def find_spikes(X, hc):
    a, b, s = hc['start'], hc['end'], hc['s']
    Y = s * X
    Q = np.arange(a, b + 1)
    ymax = Y[Q].max()
    P = [int(r) for r in Q if Y[r] == ymax]
    spikes = set(P)
    last = min(P)                                   # backward scan
    for r in range(min(P) - 1, a - 1, -1):
        between = Y[r + 1:last]
        c1 = len(between) > 0 and Y[r] > between.min() + 10
        c2 = (r == a) or Y[r] > Y[a:r].max()
        if c1 and c2:
            spikes.add(r); last = r
    last = max(P)                                   # forward scan
    for r in range(max(P) + 1, b + 1):
        between = Y[last + 1:r]
        c1 = len(between) > 0 and Y[r] > between.min() + 10
        c2 = (r == b) or Y[r] > Y[r + 1:b + 1].max()
        if c1 and c2:
            spikes.add(r); last = r
    info = []
    for r in sorted(spikes):
        primary = not (Y[r] <= ymax - 10)
        nb = [Y[j] for j in (r - 1, r + 1) if 0 <= j < len(Y)]
        if any(v == Y[r] for v in nb):
            shape = 'flat'
        elif any(v >= Y[r] - 3 for v in nb):
            shape = 'semi'
        else:
            shape = 'peaked'
        info.append({'r': r, 'primary': primary, 'shape': shape})
    nprim = sum(1 for i in info if i['primary'])
    hc.update(P=P, spikes=info, period=min(nprim, 4), ymax=ymax)
    return hc


# ----------------------------------------------------------------------------- polar phases
def polar_phases(N, S, extN, extS):
    R = len(N)
    ph = np.full(R, MAX)
    cfg = 1 if N[:30].mean() > 0 else -1          # sign of N in the current polar maximum
    start, state, min_ends = 0, MAX, []
    while start < R:
        if state == MAX:
            rN = next((r for r in range(start + 1, R) if N[r] * cfg < 0), None)
            rS = next((r for r in range(start + 1, R) if S[r] * (-cfg) < 0), None)
            if rN is None or rS is None:
                ph[start:] = MAX; break
            t = max(rN, rS); ph[start:t] = MAX; start, state = t, MIN
        elif state == MIN:
            nN = next((r for r, ty in extN if r > start and ty == -cfg), R)
            nS = next((r for r, ty in extS if r > start and ty == cfg), R)
            lN = max([r for r in range(start + 1, nN) if N[r] * cfg > 0], default=start)
            lS = max([r for r in range(start + 1, nS) if S[r] * (-cfg) > 0], default=start)
            t = max(lN, lS)
            if t >= R - 1 and (nN == R or nS == R):  # reversal still under way at the end of the data
                ph[start:] = MIN; break
            ph[start:t] = MIN; min_ends.append(t); start, state = t, REC
        else:
            t = min(R, start + STEPS_15YR); ph[start:t] = REC
            start, state, cfg = t, MAX, -cfg
    return ph, min_ends


# ----------------------------------------------------------------------------- windows
def sgn(x):
    return np.where(x > 0, 1, np.where(x < 0, -1, 0))


def crossings(X):
    s = sgn(X).astype(float)
    for i in range(1, len(s)):                      # carry sign through exact zeros
        if s[i] == 0:
            s[i] = s[i - 1]
    return [r for r in range(1, len(s)) if s[r] == -s[r - 1]]


def window_instances(dates, N, S, opts):
    """Returns list of (centre_index, half_width_days, condition_label)."""
    R = len(N)
    _, extN = longterm_extremes(N)
    _, extS = longterm_extremes(S)
    ph, min_ends = polar_phases(N, S, extN, extS)
    poles = {}
    for name, X in (('N', N), ('S', S)):
        U = minicycle_smooth(X)
        hcs = [find_spikes(X, hc) for hc in half_minicycles(X, U)]
        poles[name] = (X, hcs)
    inst = []
    for name, (X, hcs) in poles.items():
        other_X, other_hcs = poles['S' if name == 'N' else 'N']
        other_spikes = sorted(sp['r'] for h in other_hcs for sp in h['spikes'])
        for q in range(1, len(hcs) - 1):
            h, hp, hn = hcs[q], hcs[q - 1], hcs[q + 1]
            amp_ok = (abs(X[min(h['P'])] - X[min(hp['P'])]) >= 40 and
                      abs(X[min(h['P'])] - X[min(hn['P'])]) >= 40)
            if not amp_ok:
                continue
            for sp in h['spikes']:
                r = sp['r']
                if not (sgn(X[r]) == sgn(X[h['start']]) == sgn(X[h['end']])):
                    continue
                inmax = ph[r] == MAX
                pk_pr = sp['shape'] == 'peaked' and sp['primary']
                if inmax and h['period'] in (1, 2) and pk_pr:
                    inst.append((r, 20, name + ':C1 peak'))
                if inmax and abs(X[r]) > 175 and h['period'] == 1 and sp['shape'] == 'peaked':
                    inst.append((r, 25, name + ':C2 strong peak'))
                if (inmax and h['period'] == 1 and hp['period'] == 1 and hn['period'] == 1 and
                        np.abs(X[hp['start']:hp['end'] + 1]).max() > 175 and
                        np.abs(X[hn['start']:hn['end'] + 1]).max() > 175):
                    inst.append((r, 25, name + ':C3'))
                prev_o = [x for x in other_spikes if x < r]
                next_o = [x for x in other_spikes if x > r]
                if (inmax and h['period'] == 1 and prev_o and next_o and
                        abs(other_X[prev_o[-1]]) > 175 and abs(other_X[next_o[0]]) > 175):
                    inst.append((r, 25, name + ':C4'))
                if inmax and h['period'] == 3 and pk_pr:
                    inst.append((r, 15, name + ':C5 triple'))
                if inmax and h['period'] >= 4 and pk_pr:
                    inst.append((r, 10, name + ':C6 quad'))
                if sp['shape'] == 'semi':
                    inst.append((r, 10, name + ':C7 semi-flat'))
                if sp['shape'] == 'flat':
                    inst.append((r, 15, name + ':C8 flat'))
                if not sp['primary']:
                    inst.append((r, 10, name + ':C9 secondary'))
        # pole reversals
        cr = crossings(X)
        for k, r in enumerate(cr):
            prev_c = cr[k - 1] if k > 0 else 0
            nxt = (cr[k + 1] - 1) if k + 1 < len(cr) else R - 1
            if np.abs(X[prev_c:r]).max() > 15 and np.abs(X[r + 1:nxt + 1]).max(initial=0) > 15:
                inst.append((r, 15, name + ':C10 pole reversal'))
            if ph[r - 1] == MAX:
                j = r - 1
                while j > 0 and ph[j - 1] == MAX:
                    j -= 1
                if not any(j < c < r for c in cr):
                    inst.append((r, 15, name + ':C11 first reversal'))
    for r in min_ends:
        inst.append((r, 15, 'C12 end of polar minimum'))
    B = (N + S) / 2.0
    cb = crossings(B)
    for k, r in enumerate(cb):
        prev_c = cb[k - 1] if k > 0 else 0
        nxt_c = cb[k + 1] if k + 1 < len(cb) else R
        big = (np.abs(B[prev_c:r]).max() > 50 and np.abs(B[r + 1:nxt_c]).max(initial=0) > 50)
        if not (ph[r] == MAX or big):
            continue
        gap = (nxt_c - 1) - r                       # min{r'' > r : sign change after r''} - r
        if opts['c13'] == 'le9' and gap > 9:
            continue
        if opts['c13'] == 'ge9' and gap < 9:
            continue
        inst.append((r, 15, 'C13 average reversal'))
    return inst, ph, poles, min_ends


def delay_events(quakes):
    ev = sorted(d for d, m, _ in quakes if m >= 8.0)
    m75 = sorted(d for d, m, _ in quakes if m >= 7.5)
    for i, d in enumerate(m75):                     # 4th M7.5+ within 100 days
        if sum(1 for e in m75[:i + 1] if (d - e).days < 100) >= 4:
            ev.append(d)
    return sorted(set(ev))


def build_coverage(dates, inst, events, day0, day1, delay=True):
    """Union of windows as a boolean array over days day0..day1 (inclusive)."""
    n = (day1 - day0).days + 1
    cov = np.zeros(n, bool)
    for r, h, lab in inst:
        c = dates[r]
        start = c - dt.timedelta(days=h)
        shift = 0
        if delay:
            prior = [e for e in events if 0 < (start - e).days <= 100]
            if prior:
                gapd = (start - prior[-1]).days
                shift = min(20, h) if gapd < 30 else 10
        a = (start - day0).days + shift
        b = a + 2 * h
        lo, hi = max(a, 0), min(b, n - 1)
        if lo <= hi:
            cov[lo:hi + 1] = True
    return cov


def hits(cov, day0, quakes):
    out = []
    for d, m, p in quakes:
        k = (d - day0).days
        if 0 <= k < len(cov):
            out.append((d, m, p, bool(cov[k])))
    return out