"""Fill missing heights in FSL files written by convert_igra_to_fsl.py.

IGRA carries geopotential height only at mandatory levels, and the converter
wrote 9999 elsewhere. AERMET's FSL reader has no missing code for height:
9999 is read as 9999 m, which is beyond its 5000 m ceiling, so extraction of
the sounding stops at that level and AERMET sees two or three levels instead
of dozens. Its convective mixing heights then run to the 4000 m cap.

This rewrites each sounding with heights from the hypsometric equation
(pressure and temperature, integrating up from the station elevation and
re-anchoring on every real height), the same routine the server-side
converter now uses.

    python3 fix_fsl_heights.py /var/www/files/upper/2024            # report only
    python3 fix_fsl_heights.py /var/www/files/upper/2024 --write    # rewrite, originals in <dir>/.pre-heightfix/
"""

from __future__ import annotations

import math
import shutil
import sys
from pathlib import Path

G = 9.80665
RD = 287.05


def fill_heights(rows: list[list[str]], elev: float) -> tuple[list[list[str]], int]:
    n = len(rows)
    def num(s, scale, positive=False):
        """Field value, or None for the IGRA/FSL missing codes (-9999, -8888, 9999x)."""
        try:
            v = float(s)
        except ValueError:
            return None
        if v <= -8000 or v >= 99999 or (positive and v <= 0):
            return None
        return v / scale
    p = [num(r[0], 10, positive=True) for r in rows]
    z = [None if r[1] == "9999" else num(r[1], 1) for r in rows]
    t = [(v + 273.15) if (v := num(r[2], 10)) is not None else None for r in rows]
    if n and z[0] is None:
        z[0] = float(elev)
    filled = 0
    last_i = 0
    for i in range(1, n):
        if z[i] is not None:
            last_i = i
            continue
        if p[i] is None or p[last_i] is None or p[i] >= p[last_i]:
            continue
        t_here, t_last = t[i], t[last_i]
        if t_here is None and t_last is None:
            continue
        tm = t_here if t_last is None else t_last if t_here is None else 0.5 * (t_here + t_last)
        z[i] = z[last_i] + RD * tm / G * math.log(p[last_i] / p[i])
        last_i = i
        filled += 1
    known = [i for i in range(n) if z[i] is not None and p[i] is not None]
    for i in range(n):
        if z[i] is None and p[i] is not None and len(known) >= 2:
            lo = max((k for k in known if k < i), default=None)
            hi = min((k for k in known if k > i), default=None)
            if lo is not None and hi is not None and p[lo] > p[hi]:
                f = math.log(p[lo] / p[i]) / math.log(p[lo] / p[hi])
                z[i] = z[lo] + f * (z[hi] - z[lo])
                filled += 1
    out = []
    for r, zz in zip(rows, z):
        rr = list(r)
        if zz is not None and r[1] == "9999":
            rr[1] = str(int(round(zz)))
        out.append(rr)
    return out, filled


def fix_text(text: str) -> tuple[str, int, int]:
    """Returns (new text, levels filled, soundings touched)."""
    out_lines: list[str] = []
    block: list[str] = []      # lines of the current sounding
    filled = touched = 0

    def flush() -> None:
        nonlocal filled, touched
        if not block:
            return
        head = [l for l in block if l.split() and l.split()[0] in ("254", "1", "2", "3")]
        data = [l for l in block if l.split() and l.split()[0] in ("4", "5", "6", "7", "8", "9")]
        elev = 0.0
        for l in head:
            p = l.split()
            if p[0] == "1" and len(p) >= 6:
                try:
                    elev = float(p[-2])
                except ValueError:
                    pass
        rows = [l.split()[1:7] for l in data]
        if any(r[1] == "9999" for r in rows if len(r) >= 2):
            new_rows, n = fill_heights(rows, elev)
            if n:
                filled += n
                touched += 1
                out_lines.extend(head)
                for l, r in zip(data, new_rows):
                    typ = l.split()[0]
                    out_lines.append(f"{typ:>7}{r[0]:>7}{r[1]:>7}{r[2]:>7}{r[3]:>7}{r[4]:>7}{r[5]:>7}")
                block.clear()
                return
        out_lines.extend(block)
        block.clear()

    for line in text.splitlines():
        parts = line.split()
        if parts and parts[0] == "254" and len(parts) >= 5:
            flush()
        block.append(line)
    flush()
    return "\n".join(out_lines) + "\n", filled, touched


def main() -> None:
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    folder = Path(sys.argv[1])
    write = "--write" in sys.argv
    backup = folder / ".pre-heightfix"
    n_files = n_fixed = 0
    for p in sorted(folder.iterdir()):
        if not p.is_file() or not p.name.strip().upper().endswith(".FSL"):
            continue
        n_files += 1
        text = p.read_text(errors="replace")
        new, filled, touched = fix_text(text)
        if not filled:
            continue
        n_fixed += 1
        print(f"{p.name.strip():>20}  {filled:>7} heights filled in {touched:>4} soundings  {'rewritten' if write else 'would rewrite'}")
        if write:
            backup.mkdir(exist_ok=True)
            shutil.copy2(p, backup / p.name)
            tmp = p.with_name(p.name + ".tmp")
            tmp.write_text(new)
            tmp.replace(p)
    print(f"{n_files} files, {n_fixed} with missing heights{' (rewritten, originals in ' + str(backup) + ')' if write and n_fixed else ''}")


if __name__ == "__main__":
    main()