# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import zipfile
from collections import OrderedDict
from collections.abc import Generator, Iterator

import numpy as np
import torch

from earth2studio.data import IFS
from earth2studio.data.utils import fetch_data
from earth2studio.models.auto import AutoModelMixin, Package
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.models.px.base import PrognosticModel
from earth2studio.models.px.utils import PrognosticMixin
from earth2studio.utils import handshake_coords, handshake_dim
from earth2studio.utils.imports import (
    OptionalDependencyFailure,
    check_optional_dependencies,
)
from earth2studio.utils.type import CoordSystem

try:
    import anemoi.models  # noqa: F401
    import earthkit.regrid  # noqa: F401
    import ecmwf.opendata  # noqa: F401
    import flash_attn  # noqa: F401
except ImportError:
    OptionalDependencyFailure("aifsens")

VARIABLES = [
    "q50",
    "q100",
    "q150",
    "q200",
    "q250",
    "q300",
    "q400",
    "q500",
    "q600",
    "q700",
    "q850",
    "q925",
    "q1000",
    "t50",
    "t100",
    "t150",
    "t200",
    "t250",
    "t300",
    "t400",
    "t500",
    "t600",
    "t700",
    "t850",
    "t925",
    "t1000",
    "u50",
    "u100",
    "u150",
    "u200",
    "u250",
    "u300",
    "u400",
    "u500",
    "u600",
    "u700",
    "u850",
    "u925",
    "u1000",
    "v50",
    "v100",
    "v150",
    "v200",
    "v250",
    "v300",
    "v400",
    "v500",
    "v600",
    "v700",
    "v850",
    "v925",
    "v1000",
    "w50",
    "w100",
    "w150",
    "w200",
    "w250",
    "w300",
    "w400",
    "w500",
    "w600",
    "w700",
    "w850",
    "w925",
    "w1000",
    "z50",
    "z100",
    "z150",
    "z200",
    "z250",
    "z300",
    "z400",
    "z500",
    "z600",
    "z700",
    "z850",
    "z925",
    "z1000",
    "u10m",
    "v10m",
    "d2m",
    "t2m",
    "lsm",
    "msl",
    "sdor",
    "skt",
    "slor",
    "sp",
    "tcw",
    "z",
    "cp06",
    "tp06",
    "cos_latitude",
    "cos_longitude",
    "sin_latitude",
    "sin_longitude",
    "cos_julian_day",
    "cos_local_time",
    "sin_julian_day",
    "sin_local_time",
    "insolation",  # cosine zenith angle
    "stl1",
    "stl2",
    "ssrd06",
    "strd06",
    "sf",
    "tcc",
    "mcc",
    "hcc",
    "lcc",
    "u100m",
    "v100m",
    "ro",
]  # from config.json >> dataset.variables


@check_optional_dependencies()
class AIFSENS(torch.nn.Module, AutoModelMixin, PrognosticMixin):
    """Artificial Intelligence Forecasting System Ensemble (AIFS ENS v1.0), a
    probabilistic, ensemble-based forecast model from the European Centre for
    Medium-Range Weather Forecasts (ECMWF). AIFS ENS uses a GNN encoder/decoder with a
    sliding-window transformer processor, trained on ERA5 reanalysis and operational NWP
    analyses, and is run four times daily with a 6-hour time step. The model is trained
    with a CRPS objective over a small ensemble to provide calibrated probabilistic
    output.

    Note
    ----
    This model uses the checkpoints provided by ECMWF.
    For additional information see the following resources:

    - https://arxiv.org/abs/2406.01465
    - https://huggingface.co/ecmwf/aifs-ens-1.0
    - https://github.com/ecmwf/anemoi-core

    Parameters
    ----------
    model : torch.nn.Module
        Core PyTorch module with the pretrained AIFSENS weights loaded.
    latitudes : torch.Tensor
        Latitude values for the native model grid, registered as a buffer for
        interpolation.
    longitudes : torch.Tensor
        Longitude values for the native model grid, registered as a buffer for
        interpolation.
    interpolation_matrix : torch.Tensor
        CSR sparse matrix mapping ERA5 lat/lon inputs onto the native model grid.
    inverse_interpolation_matrix : torch.Tensor
        CSR sparse matrix mapping outputs from the native model grid back to ERA5
        lat/lon.
    invariants : torch.Tensor
        Tensor of shape [4, 721, 1440] containing the invariant fields "lsm", "sdor",
        "slor" and "z"

    Warning
    -------
    We encourage users to familiarize themselves with the license restrictions of this
    model's checkpoints.

    """

    VARIABLE_INVARIANTS = ["lsm", "sdor", "slor", "z"]
    VARIABLE_FORCINGS = [
        "cos_latitude",
        "cos_longitude",
        "sin_latitude",
        "sin_longitude",
        "cos_julian_day",
        "cos_local_time",
        "sin_julian_day",
        "sin_local_time",
        "insolation",
    ]

    def __init__(
        self,
        model: torch.nn.Module,
        latitudes: torch.Tensor,
        longitudes: torch.Tensor,
        interpolation_matrix: torch.Tensor,
        inverse_interpolation_matrix: torch.Tensor,
        invariants: torch.Tensor,
    ) -> None:
        super().__init__()
        self.model = model
        self.register_buffer("invariants", invariants)
        self.register_buffer("latitudes", latitudes)
        self.register_buffer("longitudes", longitudes)
        self.register_buffer("interpolation_matrix", interpolation_matrix)
        self.register_buffer(
            "inverse_interpolation_matrix", inverse_interpolation_matrix
        )

        self.invariant_ids = self.model.data_indices.data.output.forcing[
            : len(self.VARIABLE_INVARIANTS)
        ]
        self.forcing_ids = self.model.data_indices.data.output.forcing[
            len(self.VARIABLE_INVARIANTS) :
        ]

    @property
    def input_variables(self) -> list[str]:
        indices = torch.cat(
            [
                self.model.data_indices.data.input.prognostic,
                self.model.data_indices.data.input.forcing,
            ]
        )

        # Sort the concatenated tensor
        indices = indices.sort().values

        # Keep only elements NOT forcings or invariants
        mask = ~torch.isin(indices, self.model.data_indices.data.input.forcing)

        selected = [VARIABLES[i] for i in indices[mask].tolist()]
        return selected

    @property
    def output_variables(self) -> list[str]:
        # Output constants + prognostic and diagnostic - generated forcings
        indices = torch.cat(
            [
                self.model.data_indices.data.output.forcing,
                self.model.data_indices.data.output.full,
            ]
        )

        # Sort the concatenated tensor
        indices = torch.unique(indices.sort().values)

        # Keep only elements NOT in to_remove
        mask = ~torch.isin(indices, self.model.data_indices.data.output.forcing)

        selected = [VARIABLES[i] for i in indices[mask].tolist()]
        return selected

    def input_coords(self) -> CoordSystem:
        """Input coordinate system of the prognostic model
        Returns
        -------
        CoordSystem
            Coordinate system dictionary
        """
        return OrderedDict(
            {
                "batch": np.empty(0),
                "time": np.empty(0),
                "lead_time": np.array(
                    [np.timedelta64(-6, "h"), np.timedelta64(0, "h")]
                ),
                "variable": np.array(self.input_variables),
                "lat": np.linspace(90.0, -90.0, 721),
                "lon": np.linspace(0, 360, 1440, endpoint=False),
            }
        )

    @batch_coords()
    def output_coords(self, input_coords: CoordSystem) -> CoordSystem:
        """Output coordinate system of the prognostic model
        Parameters
        ----------
        input_coords : CoordSystem
            Input coordinate system to transform into output_coords
            by default None, will use self.input_coords.
        Returns
        -------
        CoordSystem
            Coordinate system dictionary
        """
        output_coords = OrderedDict(
            {
                "batch": np.empty(0),
                "time": np.empty(0),
                "lead_time": np.array([np.timedelta64(6, "h")]),
                "variable": np.array(self.output_variables),
                "lat": np.linspace(90.0, -90.0, 721),
                "lon": np.linspace(0, 360, 1440, endpoint=False),
            }
        )
        if input_coords is None:
            return output_coords

        test_coords = input_coords.copy()
        test_coords["lead_time"] = (
            test_coords["lead_time"] - input_coords["lead_time"][-1]
        )
        target_input_coords = self.input_coords()
        for i, key in enumerate(target_input_coords):
            if key not in ["batch", "time"]:
                handshake_dim(test_coords, key, i)
                handshake_coords(test_coords, target_input_coords, key)

        output_coords["batch"] = input_coords["batch"]
        output_coords["time"] = input_coords["time"]

        output_coords["lead_time"] = (
            input_coords["lead_time"][-1] + output_coords["lead_time"]
        )
        return output_coords

    @classmethod
    def load_default_package(cls) -> Package:
        """Load prognostic package"""
        package = Package(
            "hf://ecmwf/aifs-ens-1.0",
            cache_options={
                "cache_storage": Package.default_cache("aifs-ens-1.0"),
                "same_names": True,
            },
        )

        return package

    @classmethod
    @check_optional_dependencies()
    def load_model(cls, package: Package) -> PrognosticModel:
        """Load prognostic from package"""

        # Load model
        model_path = package.resolve("aifs-ens-crps-1.0.ckpt")
        model = torch.load(
            model_path, weights_only=False, map_location=torch.ones(1).device
        )
        model.eval()

        # Define the path to the metadata file
        metadata_path = "inference-anemoi-by_epoch-epoch_001-step_000040_tp_fix_0.05/anemoi-metadata/ai-models.json"

        # Extract metadata and supporting arrays from the zip file
        with zipfile.ZipFile(model_path, "r") as zipf:  # NOTE: this is totally baffling
            # Load metadata
            metadata = json.load(zipf.open(metadata_path))

            # Load supporting arrays
            supporting_arrays = {}
            for key, entry in metadata.get("supporting_arrays_paths", {}).items():
                supporting_arrays[key] = np.frombuffer(
                    zipf.read(entry["path"]),
                    dtype=entry["dtype"],
                ).reshape(entry["shape"])

        # Load interpolation matrix
        # TODO: Maybe change this to allow for multiple packages?
        interpolation_package = Package(
            "https://get.ecmwf.int/repository/earthkit/regrid/db/1/mir_16_linear",
            cache_options={
                "cache_storage": Package.default_cache(
                    "aifs-ens-1.0/interpolation_matrix"
                ),
                "same_names": True,
            },
        )
        interpolation_matrix_path = interpolation_package.resolve(
            "9533e90f8433424400ab53c7fafc87ba1a04453093311c0b5bd0b35fedc1fb83.npz"
        )
        interpolation_matrix = np.load(interpolation_matrix_path)
        torch_interpolation_matrix = torch.sparse_csr_tensor(
            crow_indices=torch.from_numpy(interpolation_matrix["indptr"]),
            col_indices=torch.from_numpy(interpolation_matrix["indices"]),
            values=torch.from_numpy(interpolation_matrix["data"]),
            size=(interpolation_matrix["shape"][0], interpolation_matrix["shape"][1]),
            dtype=torch.float64,
        )
        inverse_interpolation_package = Package(
            "https://get.ecmwf.int/repository/earthkit/regrid/db/1/mir_16_linear/",
            cache_options={
                "cache_storage": Package.default_cache(
                    "aifs-ens-1.0/inverse_interpolation_matrix"
                ),
                "same_names": True,
            },
        )
        inverse_interpolation_matrix_path = inverse_interpolation_package.resolve(
            "7f0be51c7c1f522592c7639e0d3f95bcbff8a044292aa281c1e73b842736d9bf.npz"
        )
        inverse_interpolation_matrix = np.load(inverse_interpolation_matrix_path)
        torch_inverse_interpolation_matrix = torch.sparse_csr_tensor(
            crow_indices=torch.from_numpy(inverse_interpolation_matrix["indptr"]),
            col_indices=torch.from_numpy(inverse_interpolation_matrix["indices"]),
            values=torch.from_numpy(inverse_interpolation_matrix["data"]),
            size=(
                inverse_interpolation_matrix["shape"][0],
                inverse_interpolation_matrix["shape"][1],
            ),
            dtype=torch.float64,
        )

        # Fetch invariants from IFS, note that there are deviations between these
        # invariant fields depending on where and what time the data is fetched.
        # For this model, we will use ECMWF's own invarints in the IFS data store.
        ifs = IFS(cache=True, verbose=False)
        invariants, _ = fetch_data(
            source=ifs,
            time=np.array([np.datetime64("2026-01-01T00:00:00")]),
            variable=["lsm", "sdor", "slor", "z"],
        )
        invariants = invariants.squeeze()

        # Can also fetch from NCAR ERA5 backup but these have some differences
        # invariant_package = Package(
        #     "https://nsf-ncar-era5.s3.amazonaws.com/e5.oper.invariant/197901/",
        #     cache_options={
        #         "cache_storage": Package.default_cache(
        #             "aifs-ens-1.0"
        #         ),
        #         "same_names": True,
        #     },
        # )
        # invariant_arrays = []
        # for key, value in {"lsm": 172, "sdor": 160, "slor": 163, "z": 129}.items():
        #     ds = xr.load_dataset(invariant_package.resolve(f"e5.oper.invariant.128_{value:03d}_{key}.ll025sc.1979010100_1979010100.nc"))
        #     invariant_arrays.append(ds[key.upper()].values)
        # invariants = torch.Tensor(invariant_arrays).squeeze()

        return cls(
            model,
            latitudes=torch.Tensor(supporting_arrays["latitudes"]).reshape(1, 1, -1, 1),
            longitudes=torch.Tensor(supporting_arrays["longitudes"]).reshape(
                1, 1, -1, 1
            ),
            interpolation_matrix=torch_interpolation_matrix,
            inverse_interpolation_matrix=torch_inverse_interpolation_matrix,
            invariants=invariants,
        )

    def get_cos_sin_julian_day(
        self,
        time_array: np.datetime64,
        longitudes: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Get cosine and sine of Julian day"""
        days = (
            time_array.astype("datetime64[D]") - time_array.astype("datetime64[Y]")
        ).astype(np.float32)
        hours = (
            time_array.astype("datetime64[h]") - time_array.astype("datetime64[D]")
        ).astype(np.float32)
        julian_days = days + (hours / 24.0)
        normalized = 2 * np.pi * (julian_days / 365.25)
        cos_julian_day = torch.full_like(
            longitudes, np.cos(normalized), dtype=torch.float32
        )
        sin_julian_day = torch.full_like(
            longitudes, np.sin(normalized), dtype=torch.float32
        )
        return cos_julian_day, sin_julian_day

    def get_cos_sin_local_time(
        self,
        time_array: np.datetime64,
        longitudes: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Get cosine and sine of local time"""
        hours = (
            time_array.astype("datetime64[h]") - time_array.astype("datetime64[D]")
        ).astype(np.float32)
        normalized_time = 2 * np.pi * (hours / 24.0)
        normalized_longitudes = 2 * np.pi * (longitudes / 360.0)
        tau = normalized_time + normalized_longitudes
        cos_local_time = torch.cos(tau)
        sin_local_time = torch.sin(tau)
        return cos_local_time, sin_local_time

    def get_cosine_zenith_fields(
        self,
        date: np.datetime64,
        latitudes: torch.Tensor,
        longitudes: torch.Tensor,
    ) -> torch.Tensor:
        """Get cosine zenith fields for input time array"""

        # Get Julian day
        days = (date.astype("datetime64[D]") - date.astype("datetime64[Y]")).astype(
            np.float32
        )
        hours = (date.astype("datetime64[h]") - date.astype("datetime64[D]")).astype(
            np.float32
        )
        seconds = (date.astype("datetime64[s]") - date.astype("datetime64[h]")).astype(
            np.float32
        )
        julian_day = days + seconds / 86400.0

        # Convert angle to tensor
        angle = torch.tensor(
            julian_day / 365.25 * torch.pi * 2, device=latitudes.device
        )

        # declination in [degrees]
        declination = (
            0.396372
            - 22.91327 * torch.cos(angle)
            + 4.025430 * torch.sin(angle)
            - 0.387205 * torch.cos(2 * angle)
            + 0.051967 * torch.sin(2 * angle)
            - 0.154527 * torch.cos(3 * angle)
            + 0.084798 * torch.sin(3 * angle)
        )

        # time correction in [h.degrees]
        time_correction = (
            0.004297
            + 0.107029 * torch.cos(angle)
            - 1.837877 * torch.sin(angle)
            - 0.837378 * torch.cos(2 * angle)
            - 2.340475 * torch.sin(2 * angle)
        )

        # Convert to radians
        declination = torch.deg2rad(declination)
        latitudes = torch.deg2rad(latitudes)

        # Calculate sine and cosine of declination and latitude
        sindec_sinlat = torch.sin(declination) * torch.sin(latitudes)
        cosdec_coslat = torch.cos(declination) * torch.cos(latitudes)

        # Solar hour angle
        solar_angle = torch.deg2rad((hours - 12) * 15 + longitudes + time_correction)
        zenith_angle = sindec_sinlat + cosdec_coslat * torch.cos(solar_angle)

        # Clip negative values
        return torch.clamp(zenith_angle, min=0.0)

    def _add_invariants(
        self,
        x: torch.Tensor,
        coords: CoordSystem,
    ) -> torch.Tensor:
        """add ['lsm', 'sdor', 'slor', 'z'] to input tensor"""

        shape = list(x.shape)
        shape[-3] += self.invariants.shape[-3]
        _x = torch.zeros(shape, device=x.device)

        all_ids = torch.arange(shape[-3])
        variable_ids = all_ids[~torch.isin(all_ids, self.invariant_ids)]
        _x[..., variable_ids, :, :] = x
        _x[..., self.invariant_ids, :, :] = self.invariants

        return _x

    def _prepare_input(
        self,
        x: torch.Tensor,
        coords: CoordSystem,
    ) -> torch.Tensor:
        """Prepare input tensor and coordinates for the AIFS ENS model."""
        # add invariants
        x = self._add_invariants(x, coords)

        # Interpolate the input tensor to the model grid
        shape = x.shape
        x = x.flatten(start_dim=4)
        x = x.flatten(end_dim=3)
        x = torch.swapaxes(x, 0, -1)
        x = x.to(dtype=torch.float64)
        x = self.interpolation_matrix @ x
        x = x.to(dtype=torch.float32)
        x = torch.swapaxes(x, 0, -1)
        x = x.reshape([shape[0] * shape[1], shape[2], shape[3], -1])
        x = torch.swapaxes(x, 2, 3)

        # Get cos, sin of latitude and longitude
        # (cos_latitude, sin_latitude, cos_longitude, sin_longitude)
        cos_latitude = torch.cos(torch.deg2rad(self.latitudes))
        sin_latitude = torch.sin(torch.deg2rad(self.latitudes))
        cos_longitude = torch.cos(torch.deg2rad(self.longitudes))
        sin_longitude = torch.sin(torch.deg2rad(self.longitudes))
        cos_latitude = torch.cat([cos_latitude, cos_latitude], dim=1)
        cos_longitude = torch.cat([cos_longitude, cos_longitude], dim=1)
        sin_latitude = torch.cat([sin_latitude, sin_latitude], dim=1)
        sin_longitude = torch.cat([sin_longitude, sin_longitude], dim=1)

        # Get cos, sin of Julian day
        cos_julian_day_0, sin_julian_day_0 = self.get_cos_sin_julian_day(
            coords["time"][0] - np.timedelta64(6, "h"), self.longitudes
        )
        cos_julian_day_1, sin_julian_day_1 = self.get_cos_sin_julian_day(
            coords["time"][0], self.longitudes
        )
        cos_julian_day = torch.cat([cos_julian_day_0, cos_julian_day_1], dim=1)
        sin_julian_day = torch.cat([sin_julian_day_0, sin_julian_day_1], dim=1)

        # Get cos, sin local time
        cos_local_time_0, sin_local_time_0 = self.get_cos_sin_local_time(
            coords["time"][0] - np.timedelta64(6, "h"), self.longitudes
        )
        cos_local_time_1, sin_local_time_1 = self.get_cos_sin_local_time(
            coords["time"][0], self.longitudes
        )
        cos_local_time = torch.cat([cos_local_time_0, cos_local_time_1], dim=1)
        sin_local_time = torch.cat([sin_local_time_0, sin_local_time_1], dim=1)

        # Get cosine zenith angle
        # Add insolation / cosine zenith angle
        cos_zenith_angle_0 = self.get_cosine_zenith_fields(
            coords["time"][0] - np.timedelta64(6, "h"), self.latitudes, self.longitudes
        )
        cos_zenith_angle_1 = self.get_cosine_zenith_fields(
            coords["time"][0], self.latitudes, self.longitudes
        )
        cos_zenith_angle = torch.cat([cos_zenith_angle_0, cos_zenith_angle_1], dim=1)

        # Combine inputs
        x = torch.cat(
            [
                x[:, :, :, :90],
                cos_latitude.repeat(shape[0] * shape[1], 1, 1, 1),
                cos_longitude.repeat(shape[0] * shape[1], 1, 1, 1),
                sin_latitude.repeat(shape[0] * shape[1], 1, 1, 1),
                sin_longitude.repeat(shape[0] * shape[1], 1, 1, 1),
                cos_julian_day.repeat(shape[0] * shape[1], 1, 1, 1),
                cos_local_time.repeat(shape[0] * shape[1], 1, 1, 1),
                sin_julian_day.repeat(shape[0] * shape[1], 1, 1, 1),
                sin_local_time.repeat(shape[0] * shape[1], 1, 1, 1),
                cos_zenith_angle.repeat(shape[0] * shape[1], 1, 1, 1),
                x[:, :, :, 90:],
            ],
            dim=3,
        )

        return x

    def _update_input(
        self,
        x: torch.Tensor,
        coords: CoordSystem,
    ) -> torch.Tensor:
        """Update time based inputs."""

        time0 = coords["time"][0] + coords["lead_time"][0]
        time1 = coords["time"][0] + coords["lead_time"][1]

        # Select only inputs
        # From AnemoiModelInterface.DataIndices
        # https://anemoi.readthedocs.io/projects/models/en/latest/modules/data_indices.html#usage-information
        x = x[..., self.model.data_indices.data.input.full]

        # Get cos, sin of Julian day
        cos_julian_day_0, sin_julian_day_0 = self.get_cos_sin_julian_day(
            time0, self.longitudes
        )
        cos_julian_day_1, sin_julian_day_1 = self.get_cos_sin_julian_day(
            time1, self.longitudes
        )
        cos_julian_day = torch.cat([cos_julian_day_0, cos_julian_day_1], dim=1)
        sin_julian_day = torch.cat([sin_julian_day_0, sin_julian_day_1], dim=1)

        # Get cos, sin local time
        cos_local_time_0, sin_local_time_0 = self.get_cos_sin_local_time(
            time0, self.longitudes
        )
        cos_local_time_1, sin_local_time_1 = self.get_cos_sin_local_time(
            time1, self.longitudes
        )
        cos_local_time = torch.cat([cos_local_time_0, cos_local_time_1], dim=1)
        sin_local_time = torch.cat([sin_local_time_0, sin_local_time_1], dim=1)

        # Get cosine zenith angle
        # Add insolation / cosine zenith angle
        cos_zenith_angle_0 = self.get_cosine_zenith_fields(
            time0, self.latitudes, self.longitudes
        )
        cos_zenith_angle_1 = self.get_cosine_zenith_fields(
            time1, self.latitudes, self.longitudes
        )
        cos_zenith_angle = torch.cat([cos_zenith_angle_0, cos_zenith_angle_1], dim=1)

        # Add terms to x
        x[:, :, :, 94:95] = cos_julian_day
        x[:, :, :, 95:96] = cos_local_time
        x[:, :, :, 96:97] = sin_julian_day
        x[:, :, :, 97:98] = sin_local_time
        x[:, :, :, 98:99] = cos_zenith_angle

        return x

    def _prepare_output(
        self,
        x: torch.Tensor,
        coords: CoordSystem,
    ) -> tuple[torch.Tensor, CoordSystem]:
        """Prepare input tensor and coordinates for the AIFS ENS model."""
        # Remove generated forcings
        all_indices = torch.arange(x.size(-1))
        keep = torch.isin(
            all_indices, self.model.data_indices.data.output.forcing, invert=True
        )
        x = x[..., keep]
        shape = x.shape

        # Interpolate the model grid to the lat lon grid
        x = x[:, 1:2]
        x = x.flatten(end_dim=1)
        x = torch.swapaxes(x, 0, 1)
        x = x.flatten(start_dim=1)
        x = x.to(dtype=torch.float64)
        x = self.inverse_interpolation_matrix @ x
        x = x.to(dtype=torch.float32)
        x = torch.reshape(x, [x.shape[0], shape[0], shape[-1]])
        x = torch.swapaxes(x, 0, 1)
        x = torch.swapaxes(x, 1, 2)
        x = torch.reshape(
            x,
            [
                coords["batch"].shape[0],
                coords["time"].shape[0],
                coords["lead_time"].shape[0],
                coords["variable"].shape[0],
                coords["lat"].shape[0],
                coords["lon"].shape[0],
            ],
        )

        return x

    @torch.inference_mode()
    def _forward(
        self,
        x: torch.Tensor,
        coords: CoordSystem,
        step: int = 1,
    ) -> tuple[torch.Tensor, CoordSystem]:
        output_coords = self.output_coords(coords)
        with torch.autocast(device_type=str(x.device), dtype=torch.float16):
            y = self.model.predict_step(x, fcstep=step)
            out = torch.empty(
                (x.shape[0], x.shape[1], x.shape[2], len(VARIABLES)),
                device=x.device,
            )
            out[:, 0, :, self.model.data_indices.data.input.full] = x[:, 1]
            out[:, 1, :, self.model.data_indices.data.output.full] = y[:, 0]
            out[:, 1, :, self.model.data_indices.data.input.forcing] = x[
                :, 1, :, self.model.data_indices.model.input.forcing
            ]

        return out, output_coords

    @batch_func()
    def __call__(
        self,
        x: torch.Tensor,
        coords: CoordSystem,
    ) -> tuple[torch.Tensor, CoordSystem]:
        """Runs prognostic model 1 step.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor
        coords : CoordSystem
            Input coordinate system

        Returns
        -------
        tuple[torch.Tensor, CoordSystem]
            Output tensor and coordinate system 6 hours in the future
        """
        _ = self.output_coords(coords)  # NOTE: Quick fix for exception handling
        x = self._prepare_input(x, coords)
        x, coords = self._forward(x, coords)
        x = self._prepare_output(x, coords)
        return x, coords

    def _fill_input(self, x: torch.Tensor, coords: CoordSystem) -> torch.Tensor:
        """
        Fill the model input tensor by selecting prognostic + forcing variables,
        while removing generated forcings (indices 92–100).
        """
        # add invariants to prognostics
        x = self._add_invariants(x, coords)

        batch, time, lead, _, height, width = x.shape

        # Prepare empty output tensor with VARIABLE dimension
        out = torch.empty(
            (batch, time, lead, len(VARIABLES), height, width),
            device=x.device,
        )

        # Collect relevant indices from model (prognostic + forcing)
        indices = (
            torch.cat(
                [
                    self.model.data_indices.data.input.prognostic,
                    self.model.data_indices.data.input.forcing,
                ]
            )
            .sort()
            .values
        )

        # Keep only valid indices (exclude generated forcings)
        valid_mask = ~torch.isin(indices, self.forcing_ids)

        # Fill tensor: copy input slices into selected variable slots
        out[:, :, 0, indices[valid_mask]] = x[0, 0, 0, ...]
        out[:, :, 1, indices[valid_mask]] = x[0, 0, 1, ...]

        # Drop generated forcing / invariants from output
        all_indices = torch.arange(len(VARIABLES))
        keep = torch.isin(
            all_indices, self.model.data_indices.data.output.forcing, invert=True
        )
        out = out[:, :, :, keep, ...]

        # Update coordinates with remaining variable names
        selected_variables = [VARIABLES[i] for i in all_indices[keep].tolist()]

        out_coords = coords.copy()
        out_coords["variable"] = np.array(selected_variables)

        return out, out_coords

    @batch_func()
    def _default_generator(
        self, x: torch.Tensor, coords: CoordSystem
    ) -> Generator[tuple[torch.Tensor, CoordSystem], None, None]:
        coords = coords.copy()

        self.output_coords(coords)
        first_out, coords_out = self._fill_input(x, coords)
        coords_out["lead_time"] = coords["lead_time"][1:]
        yield first_out[:, :, 1:], coords_out

        # Prepare input tensor
        x = self._prepare_input(x, coords)
        step = 1

        while True:
            # Front hook
            x, coords = self.front_hook(x, coords)

            # Forward is identity operator
            y, coords_out = self._forward(x, coords, step=step)

            # Prepare output tensor
            output_tensor = self._prepare_output(y, coords_out)

            # Rear hook
            output_tensor, coords_out = self.rear_hook(output_tensor, coords_out)

            # Yield output tensor
            yield output_tensor, coords_out.copy()

            # Update coordinates
            coords["lead_time"] = (
                coords["lead_time"]
                + self.output_coords(self.input_coords())["lead_time"]
            )
            # Prepare input tensor
            x = self._update_input(y, coords)
            step += 1

    def create_iterator(
        self, x: torch.Tensor, coords: CoordSystem
    ) -> Iterator[tuple[torch.Tensor, CoordSystem]]:
        """Creates a iterator which can be used to perform time-integration of the
        prognostic model. Will return the initial condition first (0th step).
        Parameters
        ----------
        x : torch.Tensor
            Input tensor
        coords : CoordSystem
            Input coordinate system
        Yields
        ------
        Iterator[tuple[torch.Tensor, CoordSystem]]
            Iterator that generates time-steps of the prognostic model container the
            output data tensor and coordinate system dictionary.
        """
        yield from self._default_generator(x, coords)
