import contextlib
import math
import pathlib
import re
import struct
import textwrap
import warnings

import numpy as np
import pandas as pd

import discretisedfield as df


def field_to_ovf(
    field, filename, representation="bin8", extend_scalar=False, save_subregions=True
):
    """Write the field to an OVF2.0 file.

    Data representation (``'bin4'``, ``'bin8'``, or ``'txt'``) is passed
    using ``representation`` argument. If ``extend_scalar=True``, a scalar
    field will be saved as a vector field. More precisely, if the value at
    a cell is X, that cell will be saved as (X, 0, 0).

    Parameters
    ----------
    filename : pathlib.Path, str

        Name with an extension of the file written.

    representation : str, optional

        Representation; ``'bin4'``, ``'bin8'``, or ``'txt'``. Defaults to
        ``'bin8'``.

    extend_scalar : bool, optional

        If ``True``, a scalar field will be saved as a vector field. More
        precisely, if the value at a cell is 3, that cell will be saved as
        (3, 0, 0). Defaults to ``False``.

    save_subregions : bool, optional

       If ``True`` and subregions are defined for the mesh the subregions will be saved
       to a json file. Defaults to ``True``.

    Example
    -------
    1. Write field to the OVF file.

    >>> import os
    >>> import discretisedfield as df
    ...
    >>> p1 = (0, 0, 0)
    >>> p2 = (10e-9, 5e-9, 3e-9)
    >>> n = (10, 5, 3)
    >>> mesh = df.Mesh(p1=p1, p2=p2, n=n)
    >>> value_fun = lambda point: (point[0], point[1], point[2])
    >>> field = df.Field(mesh, dim=3, value=value_fun)
    ...
    >>> filename = 'mytestfile.ohf'
    >>> field.write(filename, representation='bin8')  # write the file
    >>> os.path.isfile(filename)
    True
    >>> field_read = df.Field.fromfile(filename)  # read the file
    >>> field_read == field
    True
    >>> os.remove(filename)  # delete the file

    See also
    --------
    ~discretisedfield.Field.write
    field_from_ovf

    """
    filename = pathlib.Path(filename)
    write_dim = 3 if extend_scalar and field.dim == 1 else field.dim
    valueunits = " ".join([str(field.units) if field.units else "None"] * write_dim)
    if write_dim == 1:
        valuelabels = "field_x"
    elif extend_scalar:
        valuelabels = " ".join(["field_x"] * write_dim)
    else:
        valuelabels = " ".join(f"field_{c}" for c in field.components)

    if representation == "bin4":
        repr_string = "Binary 4"
    elif representation == "bin8":
        repr_string = "Binary 8"
    elif representation == "txt":
        repr_string = "Text"
    else:
        raise ValueError(f"Unknown {representation=}.")

    bheader = textwrap.dedent(
        f"""\
        # OOMMF OVF 2.0
        #
        # Segment count: 1
        #
        # Begin: Segment
        # Begin: Header
        #
        # Title: Field
        # Desc: File generated by Field class
        # meshunit: {field.mesh.attributes["unit"]}
        # meshtype: rectangular
        # xbase: {field.mesh.region.pmin[0] + field.mesh.cell[0]/2}
        # ybase: {field.mesh.region.pmin[1] + field.mesh.cell[1]/2}
        # zbase: {field.mesh.region.pmin[2] + field.mesh.cell[2]/2}
        # xnodes: {field.mesh.n[0]}
        # ynodes: {field.mesh.n[1]}
        # znodes: {field.mesh.n[2]}
        # xstepsize: {field.mesh.cell[0]}
        # ystepsize: {field.mesh.cell[1]}
        # zstepsize: {field.mesh.cell[2]}
        # xmin: {field.mesh.region.pmin[0]}
        # ymin: {field.mesh.region.pmin[1]}
        # zmin: {field.mesh.region.pmin[2]}
        # xmax: {field.mesh.region.pmax[0]}
        # ymax: {field.mesh.region.pmax[1]}
        # zmax: {field.mesh.region.pmax[2]}
        # valuedim: {write_dim}
        # valuelabels: {valuelabels}
        # valueunits: {valueunits}
        #
        # End: Header
        #
        # Begin: Data {repr_string}
        """
    ).encode("utf-8")

    bfooter = textwrap.dedent(
        f"""\
        # End: Data {repr_string}
        # End: Segement
        """
    ).encode("utf-8")

    reordered = field.array.transpose((2, 1, 0, 3))  # ovf ordering

    bin_rep = {"bin4": ("<f", 1234567.0), "bin8": ("<d", 123456789012345.0)}

    if save_subregions and field.mesh.subregions:
        field.mesh.save_subregions(filename)

    with open(filename, "wb") as f:
        f.write(bheader)

        if representation in bin_rep:
            # Add the binary checksum.
            f.write(struct.pack(*bin_rep[representation]))

            if extend_scalar:
                # remove scalar vector dimension
                reordered = reordered.reshape(list(reversed(field.mesh.n)))
                reordered = np.stack(
                    (reordered, np.zeros_like(reordered), np.zeros_like(reordered)),
                    axis=-1,
                )

            # processing in chuncks to reduce memory consumption
            chunksize = 100_000
            n_chunks = math.ceil(len(reordered.flat) / chunksize)
            for i in range(n_chunks):
                f.write(
                    np.asarray(
                        reordered.flat[i * chunksize : (i + 1) * chunksize],
                        dtype=bin_rep[representation][0],
                    ).tobytes()
                )
            f.write(b"\n")
        else:
            data = pd.DataFrame(reordered.reshape((-1, field.dim)))
            data.insert(loc=0, column="leading_space", value="")

            if extend_scalar:
                data.insert(loc=2, column="y", value=0.0)
                data.insert(loc=3, column="z", value=0.0)

            data.to_csv(f, sep=" ", header=False, index=False)

        f.write(bfooter)


def field_from_ovf(filename):
    """Read the field from an OVF file.

    Data representation (``txt``, ``bin4``, or ``bin8``) as well as the OVF
    version (OVF1.0 or OVF2.0) are extracted from the file itself.

    Parameters
    ----------
    filename : pathlib.Path, str

        Name of the file to be read.

    Returns
    -------
    discretisedfield.Field

        Field read from the file.

    Example
    -------
    1. Read a field from the OVF file.

    >>> import pathlib
    >>> import discretisedfield as df
    ...
    >>> current_path = pathlib.Path(__file__).absolute().parent
    >>> filepath = current_path / '..' / 'tests' / 'test_sample' / 'oommf-ovf2-bin8.omf'
    >>> field = df.Field.fromfile(filepath)
    >>> field
    Field(...)

    See also
    --------
    ~discretisedfield.Field.fromfile
    field_to_ovf

    """
    filename = pathlib.Path(filename)
    header = {}
    with open(filename, "rb") as f:
        # >>> READ HEADER <<<
        ovf_v2 = b"2.0" in next(f)
        for line in f:
            line = line.decode("utf-8")
            if line.startswith("# Begin: Data"):
                mode = line.split()[3]
                if mode == "Binary":
                    nbytes = int(line.split()[-1])
                break
            information = line[1:].split(":")  # remove leading `#`
            if len(information) > 1:
                key = information[0].strip()
                header[key] = information[1].strip()

        # valuedim is fixed to 3 and not in the header for OVF 1.0
        header["valuedim"] = int(header["valuedim"]) if ovf_v2 else 3

        # >>> MESH <<<
        p1 = (float(header[f"{key}min"]) for key in "xyz")
        p2 = (float(header[f"{key}max"]) for key in "xyz")
        cell = (float(header[f"{key}stepsize"]) for key in "xyz")
        mesh = df.Mesh(region=df.Region(p1=p1, p2=p2), cell=cell)

        nodes = math.prod(int(header[f"{key}nodes"]) for key in "xyz")

        # >>> READ DATA <<<
        if mode == "Binary":
            # OVF2 uses little-endian and OVF1 uses big-endian
            format = f'{"<" if ovf_v2 else ">"}{"d" if nbytes == 8 else "f"}'

            test_value = struct.unpack(format, f.read(nbytes))[0]
            check = {4: 1234567.0, 8: 123456789012345.0}
            if nbytes not in (4, 8) or test_value != check[nbytes]:
                raise ValueError(  # pragma: no cover
                    f"Cannot read file {filename}. The file seems to be in"
                    f" binary format ({nbytes} bytes) but the check value"
                    f" is not correct: Expected {check[nbytes]}, got"
                    f" {test_value}."
                )

            array = np.fromfile(
                f, count=int(nodes * header["valuedim"]), dtype=format
            ).reshape((-1, header["valuedim"]))
        else:
            array = pd.read_csv(
                f,
                sep=" ",
                header=None,
                dtype=np.float64,
                skipinitialspace=True,
                nrows=nodes,
                comment="#",
            )
            if len(array.columns) == header["valuedim"] + 1:
                # mumax3 writes trailing whitespace -> one extra column
                array.drop(array.columns[-1], axis=1, inplace=True)
            array = array.to_numpy()

    with contextlib.suppress(FileNotFoundError):
        mesh.load_subregions(filename)

    r_tuple = (*reversed(mesh.n), header["valuedim"])
    t_tuple = (2, 1, 0, 3)

    try:
        # multi-word components are surrounded by {}
        components = re.findall(r"(\w+|{[\w ]+})", header["valuelabels"])
    except KeyError:
        components = None
    else:

        def convert(comp):
            # Magnetization_x -> x
            # {Total field_x} -> x
            # {Total energy density} -> Total_energy_density
            comp = comp.split("_")[1] if "_" in comp else comp
            comp = comp.replace("{", "").replace("}", "")
            return "_".join(comp.split())

        components = [convert(c) for c in components]
        if len(components) != len(set(components)):  # components are not unique
            components = None

    try:
        unit_list = header["valueunits"].split()
    except KeyError:
        units = None
    else:
        if len(unit_list) == 0:
            units = None  # no unit in the file
        elif len(set(unit_list)) != 1:
            warnings.warn(
                f"File {filename} contains multiple units for the individual"
                f" components: {unit_list=}. This is not supported by"
                " discretisedfield. Units are set to None."
            )
            units = None
        else:
            units = unit_list[0]

    return df.Field(
        mesh,
        dim=header["valuedim"],
        value=array.reshape(r_tuple).transpose(t_tuple),
        components=components,
        units=units,
    )
