"""
This module is primarily intended to read model setup data from a model
definition file generated by WW3 (v6.07). There exists (to my knowledge) no
formal definition of a ww3 model_def file - as the format is primarily intended
to be used by ww3 internally. Here we need the information in the model_def
file so we can read restart files. The current implementation was generated
from reverse engineering from the Fortran source how we should interpret the
binary, unformatted sequential file (Fortran writing type). We do _not_ try
and read all data (some is ignored)- only an effort is made to read what we
need for restart files.

To note:
In case of changes in ww3 that change the format it may be easier to instead
of rewriting an interpreter just providing the needed information directly.
Specifically we need:
- frequencies
- directions
- depth
- relation linear index to ilat, ilon
- relation ilat, ilon to linear index
- number of "sea" or spatial points.

All of these can in principle be inferred form the way the model is set-up. We
do not _need_ to read the model_def file (though it is convinient to do this
automatically).

Pieter Smit, Aug, 2022

"""

import numpy

from roguewave.wavewatch3.grid_tools import Grid
from roguewave.wavewatch3.fortran_types import (
    FortranInt,
    FortranFloat,
    FortranCharacter,
)
from roguewave.wavewatch3.resources import Resource
from io import BytesIO
from typing import Union, Literal


def read_model_definition(
    resource: Resource, byte_order: Literal["<", ">", "="] = "<"
) -> Grid:

    """
    Read a ww3 model definition file, see comments on top of this file for
    further info.

    :param resource: an instance of a resource.
    :param byte_order: which byte order data is stored. Sofar currently uses
        little endian (<).
    :return: the grid, depth and the landmask.
    """

    def jump(stream, number, start):
        if start is None:
            _ = stream.read(number)
        else:
            stream.seek(start + number)

    stream = BytesIO(resource.read())

    fort_char = FortranCharacter(endianness=byte_order)
    fort_int = FortranInt(endianness=byte_order)
    fort_float = FortranFloat(endianness=byte_order)

    data = {}
    start_of_record_absolute = 0

    # The need to read the start of record (_sor) and end of record (_eor) is
    # a quirk of the sequential unformatted IO that ww3 uses to write this
    # file and our use of gfortran to compile. Specifically:
    #
    # https://gcc.gnu.org/
    #    onlinedocs/gfortran/File-format-of-unformatted-sequential-files.html
    #
    # Each record in the file contains a leading
    # start of record (sor) 4 byte marker indicating size N of the record,
    # the data written (N bytes), and a trailing end of record (eor) marker again
    # indicating the size of the data. No information on the type of the data is
    # provided, and this has to be known in advance. E.g. mulitple different
    # information types can (and are) stored as part of the same record.
    _sor = fort_int.unpack(stream, 1)[0]  # Fortran record opening
    data["id"] = fort_char.unpack(stream, 35)  # Line 70
    data["grid_name"] = fort_char.unpack(stream, 10)  # Line 70  # noqa: F841
    data["nlon"] = fort_int.unpack(stream, 1)[0]  # Number of longitude points
    data["nlat"] = fort_int.unpack(stream, 1)[0]  # Number of longitude points
    data["num_sea"] = fort_int.unpack(stream, 1)[0]  # Number of sea points
    data["num_dir"] = fort_int.unpack(stream, 1)[0]  # Number of directional points
    data["num_freq"] = fort_int.unpack(stream, 1)[0]  # Number of wavenumber points

    # Number of input bound points (see w3odatd.ftn, line 219) and
    # Number of files for output bound data (see w3odatd.ftn, line 220)
    # (3 int total, not used here, skipping, jumping to eor
    _ = jump(stream, _sor + 4, 0)
    _eor = fort_int.unpack(stream, 1)[0]
    assert _sor == _eor

    _sor = fort_int.unpack(stream, 1)[0]
    _ = jump(stream, _sor, None)
    _eor = fort_int.unpack(stream, 1)[0]
    assert _sor == _eor

    # assert fort_int.unpack(stream,1)[0] == _sor

    # Enter the section for W3GDAT (line 582)
    data["grid_type"], data["flagll"], data["iclose"] = fort_int.unpack(
        stream, None, unformatted_sequential=True
    )
    (
        data["longitude_stepsize"],
        data["latitude_stepsize"],
        data["longitude_start"],
        data["latitude_start"],
    ) = fort_float.unpack(stream, None, unformatted_sequential=True)

    _sor = fort_int.unpack(stream, 1)[0]
    start_of_record_absolute = stream.tell()
    dtype_float = numpy.dtype("float32").newbyteorder(byte_order)
    dtype_int = numpy.dtype("int32").newbyteorder(byte_order)

    # First the bottom grid. The bottom grid is only stored on the
    # computational grid
    data["bottom_datum"] = numpy.frombuffer(
        stream.read(data["num_sea"] * 4), count=data["num_sea"], dtype=dtype_float
    )

    # Next the mask layer. The masked layer is stored everywhere. Note that it
    # is stored with latitude as the fast axis- here we retrieve as longitude
    # as the fast axis (transpose)
    num_points = data["nlon"] * data["nlat"]
    data["mask"] = numpy.reshape(
        numpy.frombuffer(stream.read(num_points * 4), dtype=dtype_int),
        (data["nlon"], data["nlat"]),
    ).transpose()

    # Next the mapping that relates ix,iy -> ns. Subtract 1 to account for 0
    # based indexing (vs 1 in Fortran). Note that it is stored with latitude as
    # the fast axis- here we retrieve as longitude as the fast axis (transpose)
    data["to_linear_index"] = (
        numpy.reshape(
            numpy.frombuffer(stream.read(num_points * 4), dtype=dtype_int),
            (data["nlon"], data["nlat"]),
        ).transpose()
        - 1
    )

    num_points = data["num_sea"]
    data["to_point_index"] = (
        numpy.reshape(
            numpy.frombuffer(stream.read(num_points * 8), dtype=dtype_int),
            (2, num_points),
        )
        - 1
    )

    # FROM: w3gdatmd.ftn line 160
    #      TRFLAG    Int.  Public   Flag for use of transparencies
    #                                0: No sub-grid obstacles.
    #                                1: Obstructions at cell boundaries.
    #                                2: Obstructions at cell centers.
    #                                3: Like 1 with continuous ice.
    #                                4: Like 2 with continuous ice.
    data["tr_flag"] = fort_int.unpack(stream, 1)[0]
    _ = jump(stream, _sor, start_of_record_absolute)
    _eor = fort_int.unpack(stream, 1)[0]
    assert _sor == _eor

    # we jump three values associated with TRFLAG
    for ii in range(0, 3):
        _ = fort_int.unpack(stream, None, True)

    # Spectral parameters
    # For descriptions; see w3gdatmd.ftn
    num_spec_points = data["num_freq"] * data["num_dir"]

    # I do not fully understand why arrays we ignore are larger.. but we
    # ignore them
    byte_size_full = (num_spec_points + data["num_dir"]) * 4

    # ignore: "mapwn, mapth"
    _sor = fort_int.unpack(stream, 1)[0]
    _ = stream.read(byte_size_full * 2)
    data["direction_step_size"] = fort_float.unpack(stream, 1)[0]
    data["direction_degree"] = (
        numpy.frombuffer(stream.read(data["num_dir"] * 4), dtype=dtype_float)
        * 180
        / numpy.pi
    )

    # ignore: "stuff"
    _ = stream.read(byte_size_full * 5)

    #
    data["freq_mult_fac"] = fort_float.unpack(stream, 1)[0]
    data["start_frequency"] = fort_float.unpack(stream, 1)[0]

    # Note they store an extra frequency before start, and at the end, I
    # assume this is for convenience in calculating delta's- either way,
    # ignored here.
    data["frequency_hertz"] = (
        numpy.frombuffer(stream.read(data["num_freq"] * 4 + 8), dtype=dtype_float)
        / numpy.pi
        / 2
    )[1:-1]

    latitude = numpy.linspace(
        data["latitude_start"],
        data["latitude_start"] + data["latitude_stepsize"] * (data["nlat"] - 1),
        data["nlat"],
        endpoint=True,
    )

    longitude = numpy.linspace(
        data["longitude_start"],
        data["longitude_start"] + data["longitude_stepsize"] * data["nlon"],
        data["nlon"],
        endpoint=False,
    )
    depth = -data["bottom_datum"]
    grid = Grid(
        number_of_spatial_points=data["num_sea"],
        frequencies=data["frequency_hertz"],
        directions=data["direction_degree"],
        latitude=latitude,
        longitude=longitude,
        depth=depth,
        mask=data["mask"],
        _to_linear_index=data["to_linear_index"],
        _to_point_index=data["to_point_index"],
    )

    return grid


def _to_slice(val: Union[slice, int]) -> slice:
    if not isinstance(val, (slice, int)):
        raise ValueError("Only slice or int supported as index.")

    if isinstance(val, int):
        return slice(val, val + 1, 1)
    else:
        return val
