#!/usr/bin/env python3

# Merge the NPI and post-1984 AAD SCAR iceberg observation databases.
#
# - Reads NPI 1977-2010 and AAD 1984-2011 CSVs from cwd/input
# - Standardises dates, coordinates, counts and size-class fields
# - Retains NPI 5-class and AAD 7-class size schemes separately
# - Corrects zero Total values from size-bin sums where size counts are > 0
# - For NPI only, sets missing Total to 0 when all five size bins are explicitly 0
# - Leaves other missing Total values as missing
# - Calculates mean_classified_length_m from size-class counts
# - Writes merged CSV to cwd/output
#
# Install:
#   pip install pandas

from pathlib import Path

import pandas as pd


# =============================================================================
# CONFIGURATION
# =============================================================================

BASE_DIR = Path.cwd()
INPUT_DIR = BASE_DIR / "input"
OUTPUT_DIR = BASE_DIR / "output"

NPI_FILE = INPUT_DIR / "NPI_Iceberg_database_1977-2010_version_2022-05-20.csv"
AAD_FILE = INPUT_DIR / "AAD_Iceberg_database_1984-2011_version_2022-06-28.csv"

OUTPUT_FILE = OUTPUT_DIR / "NPI_AAD_iceberg_database_1977-2011_merged.csv"


OUTPUT_COLUMNS = [
    "source",
    "scheme",
    "Obs_ID",
    "Cruise-ID",
    "Vessel",
    "Obs_Date_ISO",
    "Latitude",
    "Longitude",
    "Total_Count",

    # NPI 5-class scheme
    "n10_50",
    "n50_200",
    "n200_500",
    "n500_1000",
    "n_gt1000",

    # AAD 7-class scheme
    "n25_100",
    "n100_200",
    "n200_400",
    "n400_800",
    "n800_1600",
    "n1600_3200",
    "n_gt3200",

    # Derived size metric
    "mean_classified_length_m",

    # NPI-only supplementary measurements
    "Freeboard",
    "Length",
    "Width",
    "Dim_as_average",
    "View_radius",
]


# =============================================================================
# HELPERS
# =============================================================================

def parse_iso_date(series: pd.Series) -> pd.Series:
    cleaned = series.str.replace("\\", "", regex=False).str.strip()

    dates = pd.to_datetime(
        cleaned,
        format="%Y-%m-%dT%H:%MZ",
        errors="coerce",
    )

    if dates.isna().any():
        bad = series[dates.isna()].head(10).tolist()
        raise ValueError(f"Could not parse observation dates: {bad}")

    return dates.dt.strftime("%Y-%m-%d")


def numeric(series: pd.Series, name: str, integer=False) -> pd.Series:
    values = series.replace("", pd.NA)
    converted = pd.to_numeric(values, errors="coerce")

    bad = values.notna() & converted.isna()

    if bad.any():
        examples = values[bad].head(10).tolist()
        raise ValueError(f"Invalid values in {name}: {examples}")

    return converted.astype("Int64") if integer else converted


def empty_int(index) -> pd.Series:
    return pd.Series(pd.NA, index=index, dtype="Int64")


def resolve_total_count(
    total: pd.Series,
    size_bins: pd.DataFrame,
    source: str,
    missing_all_zero_is_zero=False,
) -> pd.Series:

    result = total.copy()

    bin_sum = size_bins.fillna(0).sum(axis=1).astype("Int64")
    has_positive_bin = size_bins.fillna(0).gt(0).any(axis=1)

    correct = (
        result.eq(0)
        & has_positive_bin
    ).fillna(False)

    result.loc[correct] = bin_sum.loc[correct]

    missing_zero = pd.Series(False, index=result.index)
    if missing_all_zero_is_zero:
        missing_zero = (
            result.isna()
            & size_bins.notna().all(axis=1)
            & size_bins.eq(0).all(axis=1)
        )
        result.loc[missing_zero] = 0

    print(
        f"[{source}] Corrected {correct.sum():,} zero Total value(s) "
        "from size-bin sums"
    )
    if missing_all_zero_is_zero:
        print(
            f"[{source}] Set {missing_zero.sum():,} missing Total value(s) to 0 "
            "where all size bins were explicitly 0"
        )
    print(
        f"[{source}] {result.isna().sum():,} Total_Count value(s) remain missing"
    )

    return result


def npi_mean_classified_length(out: pd.DataFrame) -> pd.Series:
    """
    Estimated mean classified iceberg length.

    Closed-class representative lengths:
        10-50       -> 30 m
        50-200      -> 125 m
        200-500     -> 350 m
        500-1000    -> 750 m

    For >1000 m:
        - representative measured dimension = mean(Length, Width)
          where both exist, otherwise whichever exists
        - one >1000 berg: use measured dimension if available, else 1000 m
        - multiple and Dim_as_average = Y:
              use measured dimension for all >1000 bergs
        - multiple and Dim_as_average != Y:
              use measured dimension for one berg and 1000 m for the rest
        - no measurement:
              use 1000 m for all >1000 bergs
    """

    counts = out[
        ["n10_50", "n50_200", "n200_500", "n500_1000", "n_gt1000"]
    ].fillna(0)

    classified_count = counts.sum(axis=1)

    weighted_sum = (
        counts["n10_50"] * 30
        + counts["n50_200"] * 125
        + counts["n200_500"] * 350
        + counts["n500_1000"] * 750
    ).astype(float)

    length = out["Length"].astype("Float64")
    width = out["Width"].astype("Float64")

    # Zero/non-positive dimensions are not meaningful measurements.
    length = length.where(length > 0)
    width = width.where(width > 0)

    measured_dimension = pd.concat(
        [length, width],
        axis=1,
    ).mean(axis=1, skipna=True)

    large_count = counts["n_gt1000"].astype(int)
    dim_as_average = out["Dim_as_average"].fillna("").str.upper()

    large_contribution = pd.Series(
        0.0,
        index=out.index,
    )

    has_large = large_count > 0
    has_measurement = measured_dimension.notna()

    # Default conservative treatment: 1000 m per >1000 iceberg.
    large_contribution.loc[has_large] = (
        large_count.loc[has_large] * 1000.0
    )

    # One large iceberg: use its measured representative dimension.
    one_measured = (
        (large_count == 1)
        & has_measurement
    )

    large_contribution.loc[one_measured] = (
        measured_dimension.loc[one_measured]
    )

    # Multiple large icebergs, measurement explicitly represents their average.
    multiple_average = (
        (large_count > 1)
        & has_measurement
        & dim_as_average.eq("Y")
    )

    large_contribution.loc[multiple_average] = (
        measured_dimension.loc[multiple_average]
        * large_count.loc[multiple_average]
    )

    # Multiple large icebergs, measurement not representative of all:
    # apply it to one iceberg and use 1000 m for the remainder.
    multiple_single = (
        (large_count > 1)
        & has_measurement
        & ~dim_as_average.eq("Y")
    )

    large_contribution.loc[multiple_single] = (
        measured_dimension.loc[multiple_single]
        + (large_count.loc[multiple_single] - 1) * 1000.0
    )

    weighted_sum += large_contribution

    mean_length = weighted_sum / classified_count
    mean_length = mean_length.where(classified_count > 0)

    return mean_length.astype("Float64")


def aad_mean_classified_length(out: pd.DataFrame) -> pd.Series:
    """
    Estimated mean classified iceberg length.

    Representative lengths:
        25-100      -> 62.5 m
        100-200     -> 150 m
        200-400     -> 300 m
        400-800     -> 600 m
        800-1600    -> 1200 m
        1600-3200   -> 2400 m
        >3200       -> 3200 m conservative lower-bound estimate
    """

    counts = out[
        [
            "n25_100",
            "n100_200",
            "n200_400",
            "n400_800",
            "n800_1600",
            "n1600_3200",
            "n_gt3200",
        ]
    ].fillna(0)

    classified_count = counts.sum(axis=1)

    weighted_sum = (
        counts["n25_100"] * 62.5
        + counts["n100_200"] * 150
        + counts["n200_400"] * 300
        + counts["n400_800"] * 600
        + counts["n800_1600"] * 1200
        + counts["n1600_3200"] * 2400
        + counts["n_gt3200"] * 3200
    )

    mean_length = weighted_sum / classified_count
    mean_length = mean_length.where(classified_count > 0)

    return mean_length.astype("Float64")


# =============================================================================
# NPI
# =============================================================================

def read_npi() -> pd.DataFrame:
    columns = [
        "ID",
        "Cruise-ID",
        "Vessel",
        "Obs_Date_ISO",
        "Obs_Date_Excel",
        "Obs_Lat",
        "Obs_Lon",
        "Ice_conc",
        "RV",
        "Total",
        "no10-50",
        "no50-200",
        "no200-500",
        "no500-1000",
        "gt1000",
        "Freeboard",
        "Length",
        "Width",
        "Dim_as_average",
        "View_radius",
        "comments",
    ]

    df = pd.read_csv(
        NPI_FILE,
        sep=";",
        header=0,
        names=columns,
        dtype=str,
        keep_default_na=False,
        encoding="latin1",
    )

    out = pd.DataFrame(index=df.index)

    out["source"] = "NPI"
    out["scheme"] = "5-class"
    out["Obs_ID"] = df["ID"]
    out["Cruise-ID"] = df["Cruise-ID"]
    out["Vessel"] = df["Vessel"]
    out["Obs_Date_ISO"] = parse_iso_date(df["Obs_Date_ISO"])
    out["Latitude"] = numeric(df["Obs_Lat"], "Obs_Lat")
    out["Longitude"] = numeric(df["Obs_Lon"], "Obs_Lon")

    out["n10_50"] = numeric(df["no10-50"], "no10-50", integer=True)
    out["n50_200"] = numeric(df["no50-200"], "no50-200", integer=True)
    out["n200_500"] = numeric(df["no200-500"], "no200-500", integer=True)
    out["n500_1000"] = numeric(df["no500-1000"], "no500-1000", integer=True)
    out["n_gt1000"] = numeric(df["gt1000"], "gt1000", integer=True)

    npi_size_fields = [
        "n10_50",
        "n50_200",
        "n200_500",
        "n500_1000",
        "n_gt1000",
    ]

    out["Total_Count"] = resolve_total_count(
        numeric(df["Total"], "Total", integer=True),
        out[npi_size_fields],
        "NPI",
        missing_all_zero_is_zero=True,
    )

    for field in [
        "n25_100",
        "n100_200",
        "n200_400",
        "n400_800",
        "n800_1600",
        "n1600_3200",
        "n_gt3200",
    ]:
        out[field] = empty_int(df.index)

    out["Freeboard"] = numeric(df["Freeboard"], "Freeboard", integer=True)
    out["Length"] = numeric(df["Length"], "Length", integer=True)
    out["Width"] = numeric(df["Width"], "Width", integer=True)

    out["Dim_as_average"] = (
        df["Dim_as_average"]
        .str.strip()
        .str.upper()
        .replace("", pd.NA)
    )

    out["View_radius"] = numeric(
        df["View_radius"],
        "View_radius",
        integer=True,
    )

    out["mean_classified_length_m"] = npi_mean_classified_length(out)

    return out


# =============================================================================
# AAD
# =============================================================================

def read_aad() -> pd.DataFrame:
    # Published file has 17 header fields but 18 values per row.
    # The final value is the documented AAD 'type' field.
    columns = [
        "ID",
        "Cruise-ID",
        "Vessel",
        "Obs_Date_ISO",
        "Obs_Date_NZ",
        "Obs_Lat",
        "Obs_Lon",
        "Ice_conc",
        "RV",
        "Total",
        "size1",
        "size2",
        "size3",
        "size4",
        "size5",
        "size6",
        "size7",
        "type",
    ]

    df = pd.read_csv(
        AAD_FILE,
        header=0,
        names=columns,
        dtype=str,
        keep_default_na=False,
        encoding="latin1",
    )

    out = pd.DataFrame(index=df.index)

    out["source"] = "AAD"
    out["scheme"] = "7-class"
    out["Obs_ID"] = df["ID"]
    out["Cruise-ID"] = df["Cruise-ID"]
    out["Vessel"] = df["Vessel"]
    out["Obs_Date_ISO"] = parse_iso_date(df["Obs_Date_ISO"])
    out["Latitude"] = numeric(df["Obs_Lat"], "Obs_Lat")
    out["Longitude"] = numeric(df["Obs_Lon"], "Obs_Lon")

    for field in [
        "n10_50",
        "n50_200",
        "n200_500",
        "n500_1000",
        "n_gt1000",
    ]:
        out[field] = empty_int(df.index)

    out["n25_100"] = numeric(df["size1"], "size1", integer=True)
    out["n100_200"] = numeric(df["size2"], "size2", integer=True)
    out["n200_400"] = numeric(df["size3"], "size3", integer=True)
    out["n400_800"] = numeric(df["size4"], "size4", integer=True)
    out["n800_1600"] = numeric(df["size5"], "size5", integer=True)
    out["n1600_3200"] = numeric(df["size6"], "size6", integer=True)
    out["n_gt3200"] = numeric(df["size7"], "size7", integer=True)

    aad_size_fields = [
        "n25_100",
        "n100_200",
        "n200_400",
        "n400_800",
        "n800_1600",
        "n1600_3200",
        "n_gt3200",
    ]

    out["Total_Count"] = resolve_total_count(
        numeric(df["Total"], "Total", integer=True),
        out[aad_size_fields],
        "AAD",
    )

    for field in [
        "Freeboard",
        "Length",
        "Width",
        "View_radius",
    ]:
        out[field] = empty_int(df.index)

    out["Dim_as_average"] = pd.NA

    out["mean_classified_length_m"] = aad_mean_classified_length(out)

    return out


# =============================================================================
# MAIN
# =============================================================================

def main() -> None:
    for path in (NPI_FILE, AAD_FILE):
        if not path.exists():
            raise FileNotFoundError(f"Input not found:\n{path}")

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print(f"[READ] {NPI_FILE.name}")
    npi = read_npi()

    print(f"[READ] {AAD_FILE.name}")
    aad = read_aad()

    merged = pd.concat(
        [npi, aad],
        ignore_index=True,
    )[OUTPUT_COLUMNS]

    if merged["Latitude"].isna().any():
        raise ValueError("Missing Latitude values")

    if merged["Longitude"].isna().any():
        raise ValueError("Missing Longitude values")

    if not merged["Latitude"].between(-90, 90).all():
        raise ValueError("Latitude values outside -90 to 90")

    if not merged["Longitude"].between(-180, 180).all():
        raise ValueError("Longitude values outside -180 to 180")

    merged = merged.sort_values(
        ["Obs_Date_ISO", "source", "Obs_ID"],
        kind="stable",
    ).reset_index(drop=True)

    merged.to_csv(
        OUTPUT_FILE,
        index=False,
        encoding="utf-8-sig",
    )

    print(f"[NPI] {len(npi):,} records")
    print(f"[AAD] {len(aad):,} records")
    print(f"[TOTAL] {len(merged):,} records")
    print(
        f"[MISSING TOTAL] "
        f"{merged['Total_Count'].isna().sum():,} records"
    )
    print(
        f"[MEAN LENGTH] "
        f"{merged['mean_classified_length_m'].notna().sum():,} records calculated"
    )
    print(f"[OK] {OUTPUT_FILE.relative_to(BASE_DIR)}")


if __name__ == "__main__":
    main()