#!/usr/bin/env python3
"""Generate daily SFFI predictions over a user-defined date range.

The script reads one AMSR2 feature table per day, applies a trained SFFI
retrieval model, and preserves the year/month directory organization in the
output. The default file layout is

    INPUT_ROOT/{year}/{month_abbr}/
        amsr2_predict_samples_{date}.csv

and

    OUTPUT_ROOT/{year}/{month_abbr}/
        amsr2_predict_samples_{date}_pred.csv

where ``date`` is formatted as YYYYMMDD and ``month_abbr`` is a lowercase
three-letter month abbreviation.

The model is loaded once and reused for every day. Rows containing missing,
nonfinite, or nonnumeric model features are excluded from prediction. By
default, these rows remain in the daily output with a missing prediction and
``prediction_valid=False``. Use ``--drop-invalid-rows`` to write only rows
with finite predictions.

This script uses the model-loading and prediction utilities provided by
``predict_sffi_with_trained_model.py``. Keep both files in the same directory
when distributing or running the code.

Security
--------
Joblib and pickle files can execute arbitrary code during loading. Only use
model packages created by you or obtained from a trusted source. Native
XGBoost JSON or UBJ models are preferable for public distribution.

Example
-------
python batch_predict_daily_sffi.py \
    --start-date 2016-01-01 \
    --end-date 2025-12-31 \
    --input-root data/daily_amsr2_features \
    --output-root results/daily_sffi \
    --model-path models/XGB_SFFI_model.pkl \
    --prediction-column pred_SFFI \
    --drop-invalid-rows
"""

from __future__ import annotations

import argparse
import calendar
import json
import logging
import platform
import re
import sys
from dataclasses import dataclass
from datetime import date, datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Mapping, Sequence

import joblib
import numpy as np
import pandas as pd

try:
    from predict_sffi_with_trained_model import (
        ModelBundle,
        apply_postprocessing,
        calculate_sha256,
        inverse_target_transform,
        load_model_bundle,
        make_json_safe,
        optional_package_version,
        predict_valid_rows,
        prepare_numeric_features,
        read_training_metadata,
        resolve_postprocessing,
        resolve_target_transform,
    )
except ImportError as error:
    raise ImportError(
        "batch_predict_daily_sffi.py requires "
        "predict_sffi_with_trained_model.py in the same directory or on "
        "PYTHONPATH."
    ) from error


SCRIPT_NAME = "batch_predict_daily_sffi.py"
SCRIPT_VERSION = "1.0.0"

LOGGER = logging.getLogger("daily_sffi_batch_prediction")
SAFE_PREFIX = re.compile(r"^[A-Za-z0-9._-]+$")

DEFAULT_INPUT_TEMPLATE = (
    "{year}/{month_abbr}/amsr2_predict_samples_{date}.csv"
)
DEFAULT_OUTPUT_TEMPLATE = (
    "{year}/{month_abbr}/amsr2_predict_samples_{date}_pred.csv"
)

SUMMARY_COLUMNS = (
    "date",
    "status",
    "input_csv",
    "output_csv",
    "n_input_rows",
    "n_input_valid",
    "n_prediction_valid",
    "n_prediction_invalid",
    "n_output_rows",
    "error_type",
    "error_message",
)


@dataclass(frozen=True)
class BatchOutputPaths:
    """Batch-level manifest and provenance files."""

    manifest: Path
    success: Path
    failed: Path
    skipped: Path
    metadata: Path


def build_argument_parser() -> argparse.ArgumentParser:
    """Build the command-line interface."""

    parser = argparse.ArgumentParser(
        description=(
            "Apply a trained SFFI retrieval model to daily AMSR2 feature "
            "tables over a specified date range."
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--start-date",
        required=True,
        help="First prediction date in YYYY-MM-DD format.",
    )
    parser.add_argument(
        "--end-date",
        required=True,
        help="Last prediction date in YYYY-MM-DD format.",
    )
    parser.add_argument(
        "--input-root",
        type=Path,
        required=True,
        help="Root directory containing daily feature CSV files.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        required=True,
        help="Root directory for daily prediction CSV files.",
    )
    parser.add_argument(
        "--model-path",
        type=Path,
        required=True,
        help=(
            "Trusted Joblib package or native XGBoost JSON/UBJ model."
        ),
    )
    parser.add_argument(
        "--training-metadata",
        type=Path,
        default=None,
        help=(
            "Optional training metadata for a native XGBoost model or a "
            "model package without complete feature information."
        ),
    )
    parser.add_argument(
        "--input-template",
        default=DEFAULT_INPUT_TEMPLATE,
        help=(
            "Relative input path template. Available fields are year, month, "
            "month_abbr, date, and date_iso."
        ),
    )
    parser.add_argument(
        "--output-template",
        default=DEFAULT_OUTPUT_TEMPLATE,
        help=(
            "Relative output path template. Available fields are year, month, "
            "month_abbr, date, and date_iso."
        ),
    )
    parser.add_argument(
        "--summary-prefix",
        default="daily_sffi_prediction",
        help="Safe prefix for batch manifests and run metadata.",
    )
    parser.add_argument(
        "--model-key",
        default="model",
        help="Estimator key in a Joblib model package.",
    )
    parser.add_argument(
        "--features-key",
        default="feature_cols",
        help="Ordered feature-name key in a Joblib model package.",
    )
    parser.add_argument(
        "--target-column",
        default=None,
        help="Target-name override; otherwise saved model information is used.",
    )
    parser.add_argument(
        "--prediction-column",
        default=None,
        help="Prediction column name. Defaults to pred_<target_column>.",
    )
    parser.add_argument(
        "--target-transform",
        choices=("auto", "none", "log1p"),
        default="auto",
        help="Target transformation used during model fitting.",
    )
    parser.add_argument(
        "--postprocessing",
        choices=("auto", "none", "nonnegative", "unit_interval"),
        default="auto",
        help=(
            "Prediction constraint. Auto uses model information or the "
            "physical range associated with the target name."
        ),
    )
    parser.add_argument(
        "--feature-dtype",
        choices=("float32", "float64"),
        default="float32",
        help="Numeric type supplied to the estimator.",
    )
    parser.add_argument(
        "--drop-invalid-rows",
        action="store_true",
        help="Write only rows with finite model predictions.",
    )
    existing_group = parser.add_mutually_exclusive_group()
    existing_group.add_argument(
        "--overwrite",
        action="store_true",
        help="Replace existing daily outputs and batch summary files.",
    )
    existing_group.add_argument(
        "--skip-existing",
        action="store_true",
        help=(
            "Skip dates whose daily output already exists and refresh the "
            "batch summary files."
        ),
    )
    parser.add_argument(
        "--fail-fast",
        action="store_true",
        help="Stop immediately when any date cannot be processed.",
    )
    parser.add_argument(
        "--progress-interval",
        type=int,
        default=30,
        help="Log progress after this many attempted dates.",
    )
    parser.add_argument(
        "--log-level",
        choices=("DEBUG", "INFO", "WARNING", "ERROR"),
        default="INFO",
        help="Logging verbosity.",
    )
    return parser


def configure_logging(level: str) -> None:
    """Configure concise console logging."""

    logging.basicConfig(
        level=getattr(logging, level),
        format="%(levelname)s: %(message)s",
    )


def parse_iso_date(value: str, argument_name: str) -> date:
    """Parse one YYYY-MM-DD date."""

    try:
        return datetime.strptime(value, "%Y-%m-%d").date()
    except ValueError as error:
        raise ValueError(
            f"{argument_name} must use YYYY-MM-DD format: {value!r}."
        ) from error


def validate_prefix(prefix: str) -> str:
    """Validate a batch-summary filename prefix."""

    cleaned = prefix.strip()
    if not cleaned or SAFE_PREFIX.fullmatch(cleaned) is None:
        raise ValueError(
            "The summary prefix may contain only letters, numbers, periods, "
            "underscores, and hyphens."
        )
    return cleaned


def iter_dates(start_date: date, end_date: date):
    """Yield every date in an inclusive interval."""

    current = start_date
    while current <= end_date:
        yield current
        current += timedelta(days=1)


def render_daily_relative_path(template: str, current_date: date) -> Path:
    """Render and validate a relative daily path template."""

    context = {
        "year": f"{current_date.year:04d}",
        "month": f"{current_date.month:02d}",
        "month_abbr": calendar.month_abbr[current_date.month].lower(),
        "date": current_date.strftime("%Y%m%d"),
        "date_iso": current_date.isoformat(),
    }
    try:
        rendered = template.format(**context)
    except (KeyError, ValueError) as error:
        raise ValueError(
            f"Invalid daily path template {template!r}: {error}"
        ) from error

    relative_path = Path(rendered)
    if relative_path.is_absolute() or ".." in relative_path.parts:
        raise ValueError(
            "Daily path templates must produce relative paths without '..'."
        )
    if relative_path.suffix.lower() != ".csv":
        raise ValueError(
            f"Daily path template must produce a .csv file: {rendered}"
        )
    return relative_path


def resolve_daily_path(
    root: Path,
    template: str,
    current_date: date,
) -> tuple[Path, Path]:
    """Return absolute and relative paths for one date."""

    relative_path = render_daily_relative_path(template, current_date)
    absolute_path = root / relative_path
    return absolute_path, relative_path


def build_batch_output_paths(
    output_root: Path,
    prefix: str,
) -> BatchOutputPaths:
    """Build deterministic batch-level output paths."""

    return BatchOutputPaths(
        manifest=output_root / f"{prefix}_manifest.csv",
        success=output_root / f"{prefix}_success.csv",
        failed=output_root / f"{prefix}_failed.csv",
        skipped=output_root / f"{prefix}_skipped.csv",
        metadata=output_root / f"{prefix}_run_metadata.json",
    )


def protect_batch_outputs(
    paths: BatchOutputPaths,
    overwrite: bool,
) -> None:
    """Prevent unintended replacement of batch-level summaries."""

    all_paths = [
        paths.manifest,
        paths.success,
        paths.failed,
        paths.skipped,
        paths.metadata,
    ]
    existing = [path for path in all_paths if path.exists()]
    if existing and not overwrite:
        listed = "\n".join(f"  {path}" for path in existing)
        raise FileExistsError(
            "Batch summary files already exist:\n"
            f"{listed}\nUse --overwrite to replace them."
        )


def resolve_target_and_prediction_names(
    bundle: ModelBundle,
    target_override: str | None,
    prediction_override: str | None,
) -> tuple[str, str]:
    """Resolve output variable names."""

    target_column = (
        target_override.strip()
        if target_override is not None
        else bundle.target_column.strip()
    )
    if not target_column:
        raise ValueError("The resolved target column cannot be empty.")

    prediction_column = (
        prediction_override.strip()
        if prediction_override is not None
        else f"pred_{target_column}"
    )
    if not prediction_column:
        raise ValueError("The prediction column cannot be empty.")
    return target_column, prediction_column


def predict_daily_table(
    input_csv: Path,
    output_csv: Path,
    bundle: ModelBundle,
    target_column: str,
    prediction_column: str,
    target_transform: str,
    postprocessing: str,
    feature_dtype: str,
    drop_invalid_rows: bool,
    overwrite: bool,
) -> dict[str, Any]:
    """Predict one daily feature table and write its output CSV."""

    if not input_csv.is_file():
        raise FileNotFoundError(f"Daily input CSV not found: {input_csv}")
    if output_csv.resolve() == input_csv.resolve():
        raise ValueError("Daily input and output paths must differ.")
    if output_csv.exists() and not overwrite:
        raise FileExistsError(f"Daily output already exists: {output_csv}")

    table = pd.read_csv(input_csv, encoding="utf-8-sig")
    n_input_rows = int(len(table))
    if n_input_rows == 0:
        raise ValueError(f"Daily input CSV contains no rows: {input_csv}")
    for generated_column in (prediction_column, "prediction_valid"):
        if generated_column in table.columns:
            raise ValueError(
                f"Generated column already exists in daily input: "
                f"{generated_column}"
            )

    numeric_features, input_valid_mask = prepare_numeric_features(
        input_table=table,
        feature_names=bundle.feature_names,
        dtype=feature_dtype,
    )
    n_input_valid = int(input_valid_mask.sum())
    if n_input_valid == 0:
        raise ValueError(
            f"No rows contain all required model features: {input_csv}"
        )

    model_predictions, input_valid_positions = predict_valid_rows(
        bundle=bundle,
        numeric_features=numeric_features,
        input_valid_mask=input_valid_mask,
    )
    original_scale_predictions = inverse_target_transform(
        model_predictions=model_predictions,
        target_transform=target_transform,
    )
    final_predictions = apply_postprocessing(
        predictions=original_scale_predictions,
        mode=postprocessing,
    )
    finite_output = np.isfinite(final_predictions)
    prediction_positions = input_valid_positions[finite_output]

    table[prediction_column] = np.nan
    table["prediction_valid"] = False
    prediction_position = table.columns.get_loc(prediction_column)
    validity_position = table.columns.get_loc("prediction_valid")
    table.iloc[
        prediction_positions,
        prediction_position,
    ] = final_predictions[finite_output].astype(np.float32)
    table.iloc[
        prediction_positions,
        validity_position,
    ] = True

    if target_transform == "log1p":
        model_scale_column = f"{prediction_column}_model_scale"
        if model_scale_column in table.columns:
            raise ValueError(
                f"Generated column already exists in daily input: "
                f"{model_scale_column}"
            )
        table[model_scale_column] = np.nan
        model_scale_position = table.columns.get_loc(model_scale_column)
        finite_model_scale = np.isfinite(model_predictions)
        table.iloc[
            input_valid_positions[finite_model_scale],
            model_scale_position,
        ] = model_predictions[finite_model_scale].astype(np.float32)

    n_prediction_valid = int(finite_output.sum())
    n_prediction_invalid = n_input_rows - n_prediction_valid
    if n_prediction_valid == 0:
        raise ValueError(
            f"The model produced no finite daily predictions: {input_csv}"
        )

    if drop_invalid_rows:
        table = table.loc[table["prediction_valid"]].copy()
    n_output_rows = int(len(table))

    output_csv.parent.mkdir(parents=True, exist_ok=True)
    table.to_csv(
        output_csv,
        index=False,
        encoding="utf-8-sig",
        float_format="%.10g",
    )

    return {
        "n_input_rows": n_input_rows,
        "n_input_valid": n_input_valid,
        "n_prediction_valid": n_prediction_valid,
        "n_prediction_invalid": n_prediction_invalid,
        "n_output_rows": n_output_rows,
    }


def make_manifest_record(
    current_date: date,
    status: str,
    input_relative: Path,
    output_relative: Path,
    counts: Mapping[str, Any] | None = None,
    error: Exception | None = None,
) -> dict[str, Any]:
    """Build one stable manifest record."""

    counts = counts or {}
    return {
        "date": current_date.isoformat(),
        "status": status,
        "input_csv": input_relative.as_posix(),
        "output_csv": output_relative.as_posix(),
        "n_input_rows": counts.get("n_input_rows"),
        "n_input_valid": counts.get("n_input_valid"),
        "n_prediction_valid": counts.get("n_prediction_valid"),
        "n_prediction_invalid": counts.get("n_prediction_invalid"),
        "n_output_rows": counts.get("n_output_rows"),
        "error_type": type(error).__name__ if error is not None else None,
        "error_message": str(error) if error is not None else None,
    }


def records_to_frame(records: Sequence[Mapping[str, Any]]) -> pd.DataFrame:
    """Convert manifest records to a table with stable columns."""

    return pd.DataFrame(records, columns=SUMMARY_COLUMNS)


def save_batch_tables(
    records: Sequence[Mapping[str, Any]],
    paths: BatchOutputPaths,
) -> None:
    """Save the complete manifest and status-specific subsets."""

    manifest = records_to_frame(records)
    manifest.to_csv(
        paths.manifest,
        index=False,
        encoding="utf-8",
    )
    status_paths = {
        "success": paths.success,
        "failed": paths.failed,
        "skipped": paths.skipped,
    }
    for status, path in status_paths.items():
        subset = manifest.loc[manifest["status"] == status]
        subset.to_csv(path, index=False, encoding="utf-8")


def write_batch_metadata(
    output_path: Path,
    args: argparse.Namespace,
    start_date: date,
    end_date: date,
    bundle: ModelBundle,
    target_column: str,
    prediction_column: str,
    target_transform: str,
    postprocessing: str,
    records: Sequence[Mapping[str, Any]],
    output_paths: BatchOutputPaths,
) -> None:
    """Write machine-readable batch provenance."""

    status_counts = {
        status: sum(record["status"] == status for record in records)
        for status in ("success", "failed", "skipped")
    }
    successful_records = [
        record for record in records if record["status"] == "success"
    ]
    aggregate_counts = {
        field: int(
            sum(
                int(record[field])
                for record in successful_records
                if record[field] is not None
            )
        )
        for field in (
            "n_input_rows",
            "n_input_valid",
            "n_prediction_valid",
            "n_prediction_invalid",
            "n_output_rows",
        )
    }

    metadata = {
        "created_utc": datetime.now(timezone.utc).isoformat(),
        "script": {
            "name": SCRIPT_NAME,
            "version": SCRIPT_VERSION,
        },
        "date_range": {
            "start": start_date.isoformat(),
            "end": end_date.isoformat(),
            "requested_days": len(records),
        },
        "input_layout": {
            "root_name": args.input_root.name,
            "relative_template": args.input_template,
        },
        "output_layout": {
            "root_name": args.output_root.name,
            "relative_template": args.output_template,
        },
        "model": {
            "filename": args.model_path.name,
            "sha256": calculate_sha256(args.model_path),
            "format": bundle.model_format,
            "type": bundle.model_type,
        },
        "training_metadata": (
            {
                "filename": args.training_metadata.name,
                "sha256": calculate_sha256(args.training_metadata),
            }
            if args.training_metadata is not None
            else None
        ),
        "target_column": target_column,
        "prediction_column": prediction_column,
        "target_transform": target_transform,
        "prediction_postprocessing": postprocessing,
        "feature_dtype": args.feature_dtype,
        "feature_count": len(bundle.feature_names),
        "feature_columns": list(bundle.feature_names),
        "drop_invalid_rows": bool(args.drop_invalid_rows),
        "status_counts": status_counts,
        "aggregate_success_counts": aggregate_counts,
        "software": {
            "python": platform.python_version(),
            "numpy": np.__version__,
            "pandas": pd.__version__,
            "joblib": joblib.__version__,
            "xgboost": optional_package_version("xgboost"),
        },
        "outputs": {
            "manifest": output_paths.manifest.name,
            "success": output_paths.success.name,
            "failed": output_paths.failed.name,
            "skipped": output_paths.skipped.name,
            "run_metadata": output_paths.metadata.name,
        },
    }
    output_path.write_text(
        json.dumps(make_json_safe(metadata), indent=2, ensure_ascii=True),
        encoding="utf-8",
    )


def validate_path_templates(
    input_root: Path,
    output_root: Path,
    input_template: str,
    output_template: str,
    example_date: date,
) -> None:
    """Validate templates and ensure that example input/output paths differ."""

    input_path, _ = resolve_daily_path(
        input_root,
        input_template,
        example_date,
    )
    output_path, _ = resolve_daily_path(
        output_root,
        output_template,
        example_date,
    )
    if input_path.resolve() == output_path.resolve():
        raise ValueError(
            "Input and output templates resolve to the same daily file."
        )


def run_batch(args: argparse.Namespace) -> list[dict[str, Any]]:
    """Run daily prediction over the complete date interval."""

    start_date = parse_iso_date(args.start_date, "--start-date")
    end_date = parse_iso_date(args.end_date, "--end-date")
    if end_date < start_date:
        raise ValueError("--end-date must not precede --start-date.")
    if args.progress_interval < 1:
        raise ValueError("--progress-interval must be at least 1.")
    if not args.input_root.is_dir():
        raise NotADirectoryError(
            f"Input root directory not found: {args.input_root}"
        )

    prefix = validate_prefix(args.summary_prefix)
    validate_path_templates(
        input_root=args.input_root,
        output_root=args.output_root,
        input_template=args.input_template,
        output_template=args.output_template,
        example_date=start_date,
    )

    args.output_root.mkdir(parents=True, exist_ok=True)
    batch_paths = build_batch_output_paths(args.output_root, prefix)
    protect_batch_outputs(
        batch_paths,
        overwrite=(args.overwrite or args.skip_existing),
    )

    training_metadata = read_training_metadata(args.training_metadata)
    bundle = load_model_bundle(
        model_path=args.model_path,
        metadata=training_metadata,
        model_key=args.model_key,
        features_key=args.features_key,
    )
    target_column, prediction_column = (
        resolve_target_and_prediction_names(
            bundle=bundle,
            target_override=args.target_column,
            prediction_override=args.prediction_column,
        )
    )
    target_transform = resolve_target_transform(
        requested=args.target_transform,
        saved_transform=bundle.target_transform,
    )
    postprocessing = resolve_postprocessing(
        requested=args.postprocessing,
        target_column=target_column,
        saved_hint=bundle.postprocessing_hint,
    )

    records: list[dict[str, Any]] = []
    all_dates = list(iter_dates(start_date, end_date))
    for index, current_date in enumerate(all_dates, start=1):
        input_csv, input_relative = resolve_daily_path(
            args.input_root,
            args.input_template,
            current_date,
        )
        output_csv, output_relative = resolve_daily_path(
            args.output_root,
            args.output_template,
            current_date,
        )

        if output_csv.exists() and args.skip_existing:
            records.append(
                make_manifest_record(
                    current_date=current_date,
                    status="skipped",
                    input_relative=input_relative,
                    output_relative=output_relative,
                )
            )
            continue

        try:
            counts = predict_daily_table(
                input_csv=input_csv,
                output_csv=output_csv,
                bundle=bundle,
                target_column=target_column,
                prediction_column=prediction_column,
                target_transform=target_transform,
                postprocessing=postprocessing,
                feature_dtype=args.feature_dtype,
                drop_invalid_rows=args.drop_invalid_rows,
                overwrite=args.overwrite,
            )
            records.append(
                make_manifest_record(
                    current_date=current_date,
                    status="success",
                    input_relative=input_relative,
                    output_relative=output_relative,
                    counts=counts,
                )
            )
        except Exception as error:
            LOGGER.error(
                "Prediction failed for %s: %s",
                current_date.isoformat(),
                error,
            )
            records.append(
                make_manifest_record(
                    current_date=current_date,
                    status="failed",
                    input_relative=input_relative,
                    output_relative=output_relative,
                    error=error,
                )
            )
            if args.fail_fast:
                save_batch_tables(records, batch_paths)
                write_batch_metadata(
                    output_path=batch_paths.metadata,
                    args=args,
                    start_date=start_date,
                    end_date=end_date,
                    bundle=bundle,
                    target_column=target_column,
                    prediction_column=prediction_column,
                    target_transform=target_transform,
                    postprocessing=postprocessing,
                    records=records,
                    output_paths=batch_paths,
                )
                raise

        if index % args.progress_interval == 0 or index == len(all_dates):
            success_count = sum(
                record["status"] == "success" for record in records
            )
            failure_count = sum(
                record["status"] == "failed" for record in records
            )
            skipped_count = sum(
                record["status"] == "skipped" for record in records
            )
            LOGGER.info(
                "Processed %d/%d dates: %d success, %d failed, %d skipped.",
                index,
                len(all_dates),
                success_count,
                failure_count,
                skipped_count,
            )

    save_batch_tables(records, batch_paths)
    write_batch_metadata(
        output_path=batch_paths.metadata,
        args=args,
        start_date=start_date,
        end_date=end_date,
        bundle=bundle,
        target_column=target_column,
        prediction_column=prediction_column,
        target_transform=target_transform,
        postprocessing=postprocessing,
        records=records,
        output_paths=batch_paths,
    )

    success_count = sum(record["status"] == "success" for record in records)
    failure_count = sum(record["status"] == "failed" for record in records)
    skipped_count = sum(record["status"] == "skipped" for record in records)
    LOGGER.info(
        "Batch completed: %d success, %d failed, %d skipped.",
        success_count,
        failure_count,
        skipped_count,
    )
    LOGGER.info("Manifest: %s", batch_paths.manifest)
    LOGGER.info("Run metadata: %s", batch_paths.metadata)
    return records


def main() -> int:
    """Command-line entry point."""

    parser = build_argument_parser()
    args = parser.parse_args()
    configure_logging(args.log_level)
    try:
        records = run_batch(args)
    except Exception:
        LOGGER.exception("Daily SFFI batch prediction failed.")
        return 1

    return 1 if any(record["status"] == "failed" for record in records) else 0


if __name__ == "__main__":
    sys.exit(main())