#!/usr/bin/env python3
"""Generate SOLO_HET-Single-Counter.txt from Solar Orbiter EPD archive CDFs.

The data are obtained from the UAH EPD Data Portal API and ESA SOAR:

    UAH EPD Data Portal: https://espada.uah.es/epd/data/
    UAH EPD API:         https://espada.uah.es/epd/data/api
    ESA SOAR:            https://soar.esac.esa.int/soar/

For HET/BGO single-detector monitoring use the Level-1 single-counter products:

    epthet1-sc  EPT-HET1 ecliptic unit, Sun / Anti-Sun
    epthet2-sc  EPT-HET2 polar unit, South / North

By default the script writes the raw single-counter values, matching the
Zenodo-style SOLO_HET-Single-Counter.txt layout:

    # year month day hour minutes seconds ACC_TIME C1H C2H C1L C2L

Use --acc-time to choose the accumulation time in seconds, defaulting to
60-second rows. The ACC_TIME column contains the actual summed accumulation
seconds used for the row. Use --rates to write counts / ACC_TIME instead.
"""

from __future__ import annotations

import argparse
import json
import re
from datetime import datetime
from pathlib import Path
from typing import Iterable

import cdflib
import numpy as np
import pandas as pd
import requests


DEFAULT_API_BASE = "https://espada.uah.es/epd/data/api"
DEFAULT_VARIABLES = ("C1H", "C2H", "C1L", "C2L")
DEFAULT_ACC_TIME_SECONDS = 60


def positive_int(value: str) -> int:
    parsed = int(value)
    if parsed <= 0:
        raise argparse.ArgumentTypeError("must be a positive integer")
    return parsed


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Download Solar Orbiter EPD Level-1 EPT-HET single-counter CDFs "
            "and write the HET/BGO single-counter table."
        )
    )
    parser.add_argument("--start", required=True, help="Start time, e.g. 2026-01-15")
    parser.add_argument("--end", required=True, help="End time, e.g. 2026-01-29")
    parser.add_argument(
        "--product",
        default="epthet1-sc",
        help=(
            "EPD product to use: epthet1-sc for ecliptic Sun/Anti-Sun or "
            "epthet2-sc for polar South/North. SOAR-style epd-epthet1-sc is accepted."
        ),
    )
    parser.add_argument(
        "--variables",
        nargs="+",
        default=list(DEFAULT_VARIABLES),
        help="CDF counter variables to write, default: C1H C2H C1L C2L",
    )
    parser.add_argument(
        "--output",
        default="SOLO_HET-Single-Counter.txt",
        help="Output text file name.",
    )
    parser.add_argument(
        "--cache-dir",
        default="solo_het_cdf_cache",
        help="Directory for downloaded CDF files.",
    )
    parser.add_argument(
        "--api-base",
        default=DEFAULT_API_BASE,
        help="EPD Data Portal API base URL.",
    )
    parser.add_argument(
        "--descriptor",
        default=None,
        help="Optional API descriptor parameter.",
    )
    parser.add_argument(
        "--rates",
        action="store_true",
        help="Write count rates as counts / ACC_TIME instead of raw counts.",
    )
    parser.add_argument(
        "--acc-time",
        type=positive_int,
        default=DEFAULT_ACC_TIME_SECONDS,
        help=(
            "Output accumulation time in seconds. Counter values are summed "
            "inside these bins before writing; default: 60."
        ),
    )
    parser.add_argument(
        "--force-download",
        action="store_true",
        help="Download files again even if they already exist in the cache.",
    )
    parser.add_argument(
        "--no-download",
        action="store_true",
        help="Do not contact the API; use matching CDF files already in --cache-dir.",
    )
    return parser.parse_args()


def normalize_product(product: str) -> str:
    product = product.strip().replace("_", "-").lower()
    if product.startswith("epd-"):
        product = product[4:]
    return product


def parse_start(value: str) -> pd.Timestamp:
    return pd.Timestamp(value)


def parse_end(value: str) -> pd.Timestamp:
    ts = pd.Timestamp(value)
    if re.fullmatch(r"\d{4}-\d{2}-\d{2}", value.strip()):
        return ts + pd.Timedelta(days=1) - pd.Timedelta(nanoseconds=1)
    return ts


def api_date(value: str) -> str:
    return pd.Timestamp(value).strftime("%Y-%m-%d")


def extract_filenames(response_text: str) -> list[str]:
    text = response_text.strip()
    if not text:
        return []
    if text.startswith("["):
        data = json.loads(text)
        if isinstance(data, list):
            names = []
            for item in data:
                if isinstance(item, str):
                    names.append(item)
                elif isinstance(item, dict):
                    for key in ("filename", "file", "name"):
                        if key in item:
                            names.append(str(item[key]))
                            break
            return names
    return [line.strip() for line in text.splitlines() if line.strip().endswith(".cdf")]


def search_archive(
    api_base: str,
    product: str,
    start: str,
    end: str,
    descriptor: str | None,
) -> list[str]:
    params = {
        "level": "l1",
        "product": product,
        "from": api_date(start),
        "to": api_date(end),
    }
    if descriptor:
        params["descriptor"] = descriptor
    response = requests.get(f"{api_base.rstrip('/')}/search", params=params, timeout=60)
    response.raise_for_status()
    names = extract_filenames(response.text)
    if not names:
        raise RuntimeError(f"No CDF files returned by archive search: {response.url}")
    return names


def cached_files(cache_dir: Path, product: str) -> list[Path]:
    pattern = f"solo_L1_epd-{product}_*.cdf"
    return sorted(cache_dir.glob(pattern))


def download_files(
    api_base: str,
    filenames: Iterable[str],
    cache_dir: Path,
    force: bool = False,
) -> list[Path]:
    cache_dir.mkdir(parents=True, exist_ok=True)
    paths: list[Path] = []
    for filename in filenames:
        path = cache_dir / Path(filename).name
        if path.exists() and not force:
            paths.append(path)
            continue
        url = f"{api_base.rstrip('/')}/file/{Path(filename).name}"
        with requests.get(url, stream=True, timeout=120) as response:
            response.raise_for_status()
            tmp_path = path.with_suffix(path.suffix + ".part")
            with tmp_path.open("wb") as handle:
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        handle.write(chunk)
            tmp_path.replace(path)
        paths.append(path)
    return paths


def cdf_datetime(cdf: cdflib.CDF, var_name: str = "EPOCH") -> pd.DatetimeIndex:
    raw = cdf.varget(var_name)
    converted = cdflib.cdfepoch.to_datetime(raw)
    return pd.to_datetime(converted)


def read_cdf(path: Path, variables: list[str]) -> pd.DataFrame:
    cdf = cdflib.CDF(str(path))
    available = set(cdf.cdf_info().zVariables) | set(cdf.cdf_info().rVariables)
    required = {"EPOCH", "ACC_TIME", *variables}
    missing = sorted(required - available)
    if missing:
        raise KeyError(f"{path.name} is missing CDF variable(s): {', '.join(missing)}")

    frame = pd.DataFrame({"datetime": cdf_datetime(cdf)})
    acc_time = np.asarray(cdf.varget("ACC_TIME"), dtype=float).reshape(-1)
    frame["ACC_TIME"] = acc_time

    for var in variables:
        values = np.asarray(cdf.varget(var), dtype=float).reshape(-1)
        frame[var] = values
    return frame


def aggregate_table(
    data: pd.DataFrame,
    variables: list[str],
    acc_time_seconds: int,
    rates: bool,
) -> pd.DataFrame:
    frequency = f"{acc_time_seconds}s"
    data = data.copy()
    data["datetime"] = data["datetime"].dt.floor(frequency)

    columns = ["ACC_TIME", *variables]
    aggregated = data.groupby("datetime", as_index=False)[columns].sum(min_count=1)
    if rates:
        for var in variables:
            with np.errstate(divide="ignore", invalid="ignore"):
                aggregated[var] = aggregated[var] / aggregated["ACC_TIME"]
    return aggregated


def build_table(
    paths: list[Path],
    variables: list[str],
    start: pd.Timestamp,
    end: pd.Timestamp,
    rates: bool,
    acc_time_seconds: int,
) -> pd.DataFrame:
    if not paths:
        raise RuntimeError("No CDF files available for processing.")
    frames = [read_cdf(path, variables) for path in paths]
    data = pd.concat(frames, ignore_index=True)
    data = data.sort_values("datetime").drop_duplicates("datetime", keep="last")
    data = data.loc[(data["datetime"] >= start) & (data["datetime"] <= end)].copy()
    if data.empty:
        raise RuntimeError("No rows remain after applying the requested time range.")
    return aggregate_table(data, variables, acc_time_seconds, rates)


def format_value(value: object, rates: bool) -> str:
    if pd.isna(value):
        return "nan"
    if rates:
        return f"{float(value):.8g}"
    value = float(value)
    if value.is_integer():
        return f"{int(value)}"
    return f"{value:.8g}"


def format_acc_time(value: object) -> str:
    if pd.isna(value):
        return "nan"
    value = float(value)
    if value.is_integer():
        return f"{int(value)}"
    return f"{value:.8g}"


def write_output(path: Path, data: pd.DataFrame, variables: list[str], rates: bool) -> None:
    with path.open("w", encoding="utf-8") as handle:
        handle.write(
            "# year month day hour minutes seconds ACC_TIME "
            + " ".join(variables)
            + " \n"
        )
        for row in data.itertuples(index=False):
            timestamp: pd.Timestamp = row.datetime
            values = [getattr(row, var) for var in variables]
            formatted = " ".join(format_value(value, rates) for value in values)
            handle.write(
                f"{timestamp.year:04d} {timestamp.month:02d} {timestamp.day:02d} "
                f"{timestamp.hour:02d} {timestamp.minute:02d} {timestamp.second:02d} "
                f"{format_acc_time(row.ACC_TIME)} {formatted} \n"
            )


def main() -> int:
    args = parse_args()
    product = normalize_product(args.product)
    cache_dir = Path(args.cache_dir)
    start = parse_start(args.start)
    end = parse_end(args.end)

    if args.no_download:
        paths = cached_files(cache_dir, product)
        if not paths:
            raise SystemExit(f"No cached CDF files found in {cache_dir} for {product}.")
    else:
        filenames = search_archive(args.api_base, product, args.start, args.end, args.descriptor)
        paths = download_files(args.api_base, filenames, cache_dir, args.force_download)

    data = build_table(paths, args.variables, start, end, args.rates, args.acc_time)
    output = Path(args.output)
    write_output(output, data, args.variables, args.rates)

    mode = "rates" if args.rates else "counts"
    print(f"Wrote {len(data)} row(s) of {mode} to {output}")
    print(f"Output accumulation time: {args.acc_time} s")
    print(f"Product: {product}; CDF files: {len(paths)}; cache: {cache_dir}")
    print(f"Time range: {data['datetime'].min()} to {data['datetime'].max()}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())