"""
Input/output subroutines for GrowClust (converted from input_subs.f90).
"""

import numpy as np

from .stats import lexsort_strings


def get_inp(f):
    """Get next non-comment line from input file.

    Parameters
    ----------
    f : file object
        Open file handle.

    Returns
    -------
    line : str
        Next non-comment line (stripped).
    """
    while True:
        line = f.readline()
        if not line:
            raise EOFError("End of file reached")
        stripped = line.strip()
        if stripped and not stripped.startswith("*"):
            return stripped


def input_check(plongcutP, plongcutS, vpvs_factor, tt_dep1, tt_dep2, tt_dep3,
                tt_del1, tt_del2, tt_del3, rmsmax, delmax, iponly, nboot, maxboot, nx0, nd0):
    """Validate input parameters."""
    errors = []

    if plongcutP < 0 or plongcutS < 0:
        errors.append(f"Invalid ray parameters: plongcutP={plongcutP}, plongcutS={plongcutS}")

    if vpvs_factor < 0.0:
        errors.append(f"Invalid vpvs_factor: {vpvs_factor}")

    if tt_dep1 < 0.0 or tt_dep2 < tt_dep1 or tt_dep3 > tt_dep2:
        errors.append(f"Invalid TT table depths: {tt_dep1}, {tt_dep2}, {tt_dep3}")

    if tt_del1 < 0.0 or tt_del2 < tt_del1 or tt_del3 > tt_del2:
        errors.append(f"Invalid TT table distances: {tt_del1}, {tt_del2}, {tt_del3}")

    ndel = int(np.floor((tt_del2 + tt_del3 / 10.0 - tt_del1) / tt_del3)) + 1
    ndep = int(np.floor((tt_dep2 + tt_dep3 / 10.0 - tt_dep1) / tt_dep3)) + 1
    if ndel > nx0:
        errors.append(f"Too many X points in TT table: {ndel} > {nx0}")
    if ndep > nd0:
        errors.append(f"Too many Z points in TT table: {ndep} > {nd0}")

    if rmsmax <= 0.0 or delmax <= 0.0:
        errors.append(f"Invalid rmsmax={rmsmax} or delmax={delmax}")

    if iponly < 0 or iponly > 2:
        errors.append(f"Invalid iponly: {iponly}")

    if nboot < 0 or nboot > maxboot:
        errors.append(f"Invalid nboot={nboot}, maxboot={maxboot}")

    if errors:
        for e in errors:
            print(f"Input error: {e}")
        raise ValueError("Input parameter validation failed")

    print("Input parameters ok!")


def param_check(hshiftmax, vshiftmax, rmedmax, boxwid, nit, samp_type, irelonorm, vzmodel_type):
    """Validate grow_params parameters."""
    errors = []

    if hshiftmax <= 0.0 or vshiftmax <= 0.0:
        errors.append(f"Invalid hshiftmax={hshiftmax}, vshiftmax={vshiftmax}")

    if rmedmax < 0.0:
        errors.append(f"Invalid rmedmax: {rmedmax}")

    if boxwid <= 0.0 or nit >= 200:
        errors.append(f"Invalid boxwid={boxwid}, nit={nit}")

    if samp_type < 1 or samp_type > 2:
        errors.append(f"Invalid samp_type: {samp_type}")

    if irelonorm < 1 or irelonorm > 3:
        errors.append(f"Invalid irelonorm: {irelonorm}")

    if vzmodel_type < 1 or vzmodel_type > 3:
        errors.append(f"Invalid vzmodel_type: {vzmodel_type}")

    if errors:
        for e in errors:
            print(f"Param error: {e}")
        raise ValueError("Parameter validation failed")

    print("grow_params parameters ok!")


def read_evfile(ievform, evfile, maxevid):
    """Read event list to get starting (catalog) locations.

    Parameters
    ----------
    ievform : int
        Event list format: 0=evlist, 1=phase, 2=growclust_cat, 3=hypoinverse.
    evfile : str
        Event list file name.
    maxevid : int
        Maximum event ID.

    Returns
    -------
    events : dict with keys:
        idcusp, qyr_cat, qmon_cat, qdy_cat, qhr_cat, qmn_cat, qsc_cat,
        qmag_cat, qlat_cat, qlon_cat, qdep_cat, qid2qnum, nq,
        min_qdep, max_qdep, max_qid
    """
    if ievform > 3:
        raise ValueError(f"Unknown event list format: {ievform}")

    idcusp = []
    qyr_cat = []
    qmon_cat = []
    qdy_cat = []
    qhr_cat = []
    qmn_cat = []
    qsc_cat = []
    qmag_cat = []
    qlat_cat = []
    qlon_cat = []
    qdep_cat = []

    with open(evfile, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            if ievform == 0:
                # evlist format (fixed-width)
                parts = line.split()
                # Format: ID yr mo dy hr mn sec evtype mag magtype lat lon dep
                # Skip first field (10 chars), then read id
                evid = int(parts[0]) if len(parts[0]) <= 9 else int(line[10:19].strip())
                # Try parsing as free-format
                try:
                    evid = int(parts[0])
                    yr = int(parts[1])
                    mo = int(parts[2])
                    dy = int(parts[3])
                    hr = int(parts[4])
                    mn = int(parts[5])
                    sc = float(parts[6])
                    mag = float(parts[8]) if len(parts) > 8 else 0.0
                    lat = float(parts[-3])
                    lon = float(parts[-2])
                    dep = float(parts[-1])
                except (ValueError, IndexError):
                    continue

            elif ievform == 1:
                # phase format
                parts = line.split()
                if len(parts) < 14:
                    continue
                yr = int(parts[0])
                mo = int(parts[1])
                dy = int(parts[2])
                hr = int(parts[3])
                mn = int(parts[4])
                sc = float(parts[5])
                lat = float(parts[6])
                lon = float(parts[7])
                dep = float(parts[8])
                mag = float(parts[9])
                evid = int(parts[13])

            elif ievform == 2:
                # growclust_cat format
                parts = line.split()
                if len(parts) < 11:
                    continue
                yr = int(parts[0])
                mo = int(parts[1])
                dy = int(parts[2])
                hr = int(parts[3])
                mn = int(parts[4])
                sc = float(parts[5])
                evid = int(parts[6])
                lat = float(parts[7])
                lon = float(parts[8])
                dep = float(parts[9])
                mag = float(parts[10])

            elif ievform == 3:
                # hypoinverse format
                parts = line.split()
                if len(parts) < 14:
                    continue
                evid = int(parts[0])
                yr = int(parts[1])
                mo = int(parts[2])
                dy = int(parts[3])
                hr = int(parts[4])
                mn = int(parts[5])
                sc = float(parts[6])
                lat = float(parts[7])
                lon = float(parts[8])
                dep = float(parts[9])
                mag = float(parts[13])

            idcusp.append(evid)
            qyr_cat.append(yr)
            qmon_cat.append(mo)
            qdy_cat.append(dy)
            qhr_cat.append(hr)
            qmn_cat.append(mn)
            qsc_cat.append(sc)
            qmag_cat.append(mag)
            qlat_cat.append(lat)
            qlon_cat.append(lon)
            qdep_cat.append(dep)

    nq = len(idcusp)
    idcusp = np.array(idcusp, dtype=np.int64)
    qyr_cat = np.array(qyr_cat, dtype=np.int32)
    qmon_cat = np.array(qmon_cat, dtype=np.int32)
    qdy_cat = np.array(qdy_cat, dtype=np.int32)
    qhr_cat = np.array(qhr_cat, dtype=np.int32)
    qmn_cat = np.array(qmn_cat, dtype=np.int32)
    qsc_cat = np.array(qsc_cat, dtype=np.float32)
    qmag_cat = np.array(qmag_cat, dtype=np.float32)
    qlat_cat = np.array(qlat_cat, dtype=np.float64)
    qlon_cat = np.array(qlon_cat, dtype=np.float64)
    qdep_cat = np.array(qdep_cat, dtype=np.float32)

    max_qid = int(np.max(idcusp)) if nq > 0 else 0
    if max_qid > maxevid:
        raise ValueError(f"Event ID too large: {max_qid} > {maxevid}")

    # Build qid2qnum lookup
    qid2qnum = np.zeros(max_qid + 1, dtype=np.int32)
    for i, eid in enumerate(idcusp):
        qid2qnum[eid] = i + 1  # 1-based index like Fortran

    min_qdep = float(np.min(qdep_cat)) if nq > 0 else 0.0
    max_qdep = float(np.max(qdep_cat)) if nq > 0 else 0.0

    print(f"Read {nq} events, depth range: [{min_qdep:.2f}, {max_qdep:.2f}] km")

    return {
        "idcusp": idcusp,
        "qid2qnum": qid2qnum,
        "qyr_cat": qyr_cat,
        "qmon_cat": qmon_cat,
        "qdy_cat": qdy_cat,
        "qhr_cat": qhr_cat,
        "qmn_cat": qmn_cat,
        "qsc_cat": qsc_cat,
        "qmag_cat": qmag_cat,
        "qlat_cat": qlat_cat,
        "qlon_cat": qlon_cat,
        "qdep_cat": qdep_cat,
        "nq": nq,
        "min_qdep": min_qdep,
        "max_qdep": max_qdep,
        "max_qid": max_qid,
    }


def read_stlist(stfile, istform):
    """Read station list and return sorted arrays with locations.

    Parameters
    ----------
    stfile : str
        Station list file name.
    istform : int
        Station list format (0, 1, or 2).

    Returns
    -------
    stnames : list of str
        Sorted station names (5 chars).
    stlats : ndarray
        Station latitudes.
    stlons : ndarray
        Station longitudes.
    stelevs : ndarray
        Station elevations (m).
    nsta : int
        Number of stations.
    stkey : dict
        Maps first character -> starting index in sorted list.
    """
    names = []
    lats = []
    lons = []
    elevs = []

    with open(stfile, "r") as f:
        for line in f:
            line_raw = line.rstrip("\n")
            if not line_raw.strip():
                continue

            if istform == 0:
                # SEED format: skip 3 chars, read 5 char name, skip 10, read lat/lon
                name = line_raw[3:8].strip()
                parts = line_raw[18:].split()
                if len(parts) < 2:
                    continue
                lat = float(parts[0])
                lon = float(parts[1])
                elev = 0.0
            elif istform <= 2:
                parts = line_raw.split()
                if len(parts) < 3:
                    continue

                # Check for NW_STNAME format
                name = parts[0]
                if len(name) > 2 and "_" not in name:
                    pass
                elif " " in line_raw[:3] and line_raw[0] != " ":
                    # NW STNAME format
                    name = parts[1] if len(parts) > 1 else parts[0]

                # Truncate to 5 chars
                name = name[:5]

                if istform == 2 and len(parts) >= 4:
                    lat = float(parts[1]) if len(parts) > 2 else float(parts[-2])
                    lon = float(parts[2]) if len(parts) > 3 else float(parts[-1])
                    try:
                        lat = float(parts[-3])
                        lon = float(parts[-2])
                        elev = float(parts[-1])
                    except (ValueError, IndexError):
                        lat = float(parts[1])
                        lon = float(parts[2])
                        elev = float(parts[3]) if len(parts) > 3 else 0.0
                else:
                    try:
                        lat = float(parts[-2])
                        lon = float(parts[-1])
                    except (ValueError, IndexError):
                        lat = float(parts[1])
                        lon = float(parts[2])
                    elev = 0.0
            else:
                raise ValueError(f"Undefined station list type: {istform}")

            names.append(name.ljust(5)[:5])
            lats.append(lat)
            lons.append(lon)
            elevs.append(elev)

    # Sort alphabetically
    isort = lexsort_strings(names)
    sorted_names = [names[i] for i in isort]
    sorted_lats = np.array([lats[i] for i in isort], dtype=np.float64)
    sorted_lons = np.array([lons[i] for i in isort], dtype=np.float64)
    sorted_elevs = np.array([elevs[i] for i in isort], dtype=np.float32)

    # Remove duplicates
    unique_names = []
    unique_lats = []
    unique_lons = []
    unique_elevs = []
    for i in range(len(sorted_names)):
        if i == 0 or sorted_names[i] != sorted_names[i - 1]:
            unique_names.append(sorted_names[i])
            unique_lats.append(sorted_lats[i])
            unique_lons.append(sorted_lons[i])
            unique_elevs.append(sorted_elevs[i])

    nsta = len(unique_names)
    unique_lats = np.array(unique_lats, dtype=np.float64)
    unique_lons = np.array(unique_lons, dtype=np.float64)
    unique_elevs = np.array(unique_elevs, dtype=np.float32)

    # Build stkey: maps first character code -> starting index
    stkey = {}
    for k in range(nsta):
        ch = unique_names[k][0]
        if ch not in stkey:
            stkey[ch] = k

    print(f"Station locations read. Nsta = {nsta}")
    return unique_names, unique_lats, unique_lons, unique_elevs, nsta, stkey


def lookup_sta(stname, nsta, snames, skeys, slats, slons, selevs):
    """Look up station location by name from sorted station list.

    Returns
    -------
    slat, slon, selev : float
        Station location, or (-999, -999, -999) if not found.
    """
    stname5 = stname.ljust(5)[:5]
    ch = stname5[0]
    istart = skeys.get(ch, 0)

    for i in range(istart, nsta):
        if snames[i] == stname5:
            return slats[i], slons[i], selevs[i]
        # Early exit if we've passed alphabetically
        if snames[i] > stname5:
            break

    return -999.0, -999.0, -999.0


def read_xcordata(irxform, it12form, xcorfile, nq, max_qid, qid2qnum,
                  qlat, qlon, rmincut, rmin_ngood, delmax, rpsavgmin,
                  iponly, ngoodmin, nsta, slnames, sllats, sllons, slelevs, slkeys):
    """Read cross-correlation data from input file.

    Parameters
    ----------
    irxform : int
        Xcor data format: 0=binary, 1=dt.cc text.
    it12form : int
        Sign convention: 12 (t1-t2) or 21 (t2-t1).
    xcorfile : str
        Cross-correlation data file name.
    nq : int
        Total number of events.
    max_qid : int
        Maximum event ID.
    qid2qnum : ndarray
        Maps event ID to serial event number (1-based).
    qlat, qlon : ndarray
        Event latitudes and longitudes.
    rmincut : float
        Minimum rxcor to keep observation.
    rmin_ngood : float
        Minimum rxcor for "good" observation.
    delmax : float
        Maximum station distance.
    rpsavgmin : float
        Minimum avg rxcor for pair.
    iponly : int
        0=use P and S, 1=P only.
    ngoodmin : int
        Minimum "good" observations per pair.
    nsta : int
        Number of stations.
    slnames, sllats, sllons, slelevs : arrays
        Station list data.
    slkeys : dict
        Station lookup keys.

    Returns
    -------
    xcor_data : dict
        Contains: npair, nk, iqq1, iqq2, idcusp11, idcusp22,
        index1, index2, stname, ipp, tdif, rxcor, dist, slat, slon
    """
    degkm = 111.1949266
    degrad = 180.0 / np.pi

    if irxform == 0:
        raise NotImplementedError("Binary xcor format (irxform=0) not supported in Python version. Use text format (irxform=1).")

    if irxform != 1:
        raise ValueError(f"Unknown xcor data format: {irxform}")

    # Text format (dt.cc)
    print(f"Reading xcor-data text file: {xcorfile}")

    iqq1_list = []
    iqq2_list = []
    idcusp11_list = []
    idcusp22_list = []
    index1_list = []
    index2_list = []
    stname_list = []
    ipp_list = []
    tdif_list = []
    rxcor_list = []
    dist_list = []
    slat_list = []
    slon_list = []

    npair = 0
    nk = 0
    npair_cut = 0
    nk_cut = 0

    # Temporary storage per pair
    current_pair = None
    pair_data = []

    def process_pair(qnum1, qnum2, qcusp1, qcusp2, pair_data):
        nonlocal npair, nk, npair_cut, nk_cut

        if not pair_data:
            return

        kk = len(pair_data)
        rps_pr = sum(d["rxcor"] for d in pair_data) / kk

        # Determine which to keep
        keep = []
        ngood_pr = 0
        for d in pair_data:
            if iponly == 1 and d["ipp"] != 1:
                nk_cut += 1
                continue
            if d["rxcor"] < rmincut:
                nk_cut += 1
                continue
            keep.append(d)
            if d["rxcor"] >= rmin_ngood and d["dist"] <= delmax:
                ngood_pr += 1

        nkeep_pr = len(keep)

        if nkeep_pr > 0 and ngood_pr >= ngoodmin and rps_pr >= rpsavgmin:
            npair += 1
            iqq1_list.append(qnum1)
            iqq2_list.append(qnum2)
            idcusp11_list.append(qcusp1)
            idcusp22_list.append(qcusp2)
            k1 = nk
            for d in keep:
                stname_list.append(d["stname"])
                ipp_list.append(d["ipp"])
                tdif_list.append(d["tdif"])
                rxcor_list.append(d["rxcor"])
                dist_list.append(d["dist"])
                slat_list.append(d["slat"])
                slon_list.append(d["slon"])
                nk += 1
            index1_list.append(k1)
            index2_list.append(nk - 1)
        else:
            npair_cut += 1
            nk_cut += nkeep_pr

    qnum1 = 0
    qnum2 = 0
    qcusp1 = 0
    qcusp2 = 0

    with open(xcorfile, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            if line.startswith("#"):
                # Process previous pair
                if current_pair is not None:
                    process_pair(qnum1, qnum2, qcusp1, qcusp2, pair_data)

                # New pair header
                parts = line[1:].split()
                qcusp1 = int(parts[0])
                qcusp2 = int(parts[1])

                # Look up serial IDs
                qnum1 = int(qid2qnum[qcusp1])
                qnum2 = int(qid2qnum[qcusp2])

                if qnum1 < 1 or qnum2 < 1:
                    raise ValueError(f"Event pair in xcor file missing from event list: {qcusp1}, {qcusp2}")

                current_pair = (qcusp1, qcusp2)
                pair_data = []

            else:
                # Differential time line
                parts = line.split()
                if len(parts) < 4:
                    continue

                stname = parts[0].strip()[:5].ljust(5)
                tdif_val = float(parts[1])
                rxcor_val = float(parts[2])
                ippchar = parts[3]

                if it12form == 12:
                    tdif_val = -tdif_val

                if ippchar.upper().startswith("S"):
                    ipp_val = 2
                else:
                    ipp_val = 1

                # Look up station
                sta_lat, sta_lon, sta_elev = lookup_sta(
                    stname, nsta, slnames, slkeys, sllats, sllons, slelevs)

                if sta_lat < -99.0:
                    raise ValueError(f"Station not found: {stname}")

                # Compute distance to event pair centroid
                qlat0 = (qlat[qnum1 - 1] + qlat[qnum2 - 1]) / 2.0
                qlon0 = (qlon[qnum1 - 1] + qlon[qnum2 - 1]) / 2.0
                reflat0 = (sta_lat + qlat0) / 2.0
                cosreflat0 = np.cos(reflat0 / degrad)
                dx = (sta_lon - qlon0) * cosreflat0
                dy = sta_lat - qlat0
                dist_val = np.sqrt(dx**2 + dy**2) * degkm

                pair_data.append({
                    "stname": stname,
                    "tdif": tdif_val,
                    "rxcor": rxcor_val,
                    "ipp": ipp_val,
                    "dist": dist_val,
                    "slat": sta_lat,
                    "slon": sta_lon,
                })

    # Process last pair
    if current_pair is not None:
        process_pair(qnum1, qnum2, qcusp1, qcusp2, pair_data)

    print(f"npair_kept={npair}, nk_kept={nk}")
    print(f"npair_cut={npair_cut}, nk_cut={nk_cut}")

    return {
        "npair": npair,
        "nk": nk,
        "iqq1": np.array(iqq1_list, dtype=np.int32),
        "iqq2": np.array(iqq2_list, dtype=np.int32),
        "idcusp11": np.array(idcusp11_list, dtype=np.int64),
        "idcusp22": np.array(idcusp22_list, dtype=np.int64),
        "index1": np.array(index1_list, dtype=np.int32),
        "index2": np.array(index2_list, dtype=np.int32),
        "stname": stname_list,
        "ipp": np.array(ipp_list, dtype=np.int32),
        "tdif": np.array(tdif_list, dtype=np.float32),
        "rxcor": np.array(rxcor_list, dtype=np.float32),
        "dist": np.array(dist_list, dtype=np.float32),
        "slat": np.array(slat_list, dtype=np.float64),
        "slon": np.array(slon_list, dtype=np.float64),
    }
