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


class _FieldIO_OVF:
    __slots__ = []

    def _to_ovf(
        self, filename, representation="bin8", extend_scalar=False, save_subregions=True
    ):
        if self.mesh.region.ndim != 3:
            raise RuntimeError(
                "OVF files can only store fields with 'ndim=3', not"
                f" {self.mesh.region.ndim=}."
            )
        filename = pathlib.Path(filename)
        write_dim = 3 if extend_scalar and self.nvdim == 1 else self.nvdim
        valueunits = " ".join([str(self.unit) if self.unit 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 self.vdims)

        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=}.")

        if len(set(self.mesh.region.units)) > 1:
            raise ValueError(
                "OVF files only support fields with the same units in all spatial"
                " directions."
            )

        bheader = textwrap.dedent(
            f"""\
            # OOMMF OVF 2.0
            #
            # Segment count: 1
            #
            # Begin: Segment
            # Begin: Header
            #
            # Title: Field
            # Desc: File generated by Field class
            # meshunit: {self.mesh.region.units[0]}
            # meshtype: rectangular
            # xbase: {self.mesh.region.pmin[0] + self.mesh.cell[0]/2}
            # ybase: {self.mesh.region.pmin[1] + self.mesh.cell[1]/2}
            # zbase: {self.mesh.region.pmin[2] + self.mesh.cell[2]/2}
            # xnodes: {self.mesh.n[0]}
            # ynodes: {self.mesh.n[1]}
            # znodes: {self.mesh.n[2]}
            # xstepsize: {self.mesh.cell[0]}
            # ystepsize: {self.mesh.cell[1]}
            # zstepsize: {self.mesh.cell[2]}
            # xmin: {self.mesh.region.pmin[0]}
            # ymin: {self.mesh.region.pmin[1]}
            # zmin: {self.mesh.region.pmin[2]}
            # xmax: {self.mesh.region.pmax[0]}
            # ymax: {self.mesh.region.pmax[1]}
            # zmax: {self.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: Segment
            """
        ).encode("utf-8")

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

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

        if save_subregions and self.mesh.subregions:
            self.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(self.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, self.nvdim)))
                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)

    @classmethod
    def _from_ovf(cls, filename):
        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"]
            units = [header["meshunit"]] * 3
            mesh = df.Mesh(region=df.Region(p1=p1, p2=p2, units=units), 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 vdims are surrounded by {}
            vdims = re.findall(r"(\w+|{[\w ]+})", header["valuelabels"])
        except KeyError:
            vdims = 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())

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

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

        return cls(
            mesh,
            nvdim=header["valuedim"],
            value=array.reshape(r_tuple).transpose(t_tuple),
            vdims=vdims,
            unit=unit,
        )
