# ovf2io is a utility for OOMMF Vector Field (.ovf) IO developed by WSP as a member of the McMorran Lab
# Copyright (C) 2023  William S. Parker
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
import struct
import shlex
import numpy as np
from . import _templates

def _generate_coordinates(x, y, z, shape, meshtype):
    generated = False
    if x is None and y is None and z is None:
        x = np.arange(shape[0]) if meshtype == 'rectangular' else np.arange(shape[0])
        y = np.arange(shape[1]) if meshtype == 'rectangular' else np.zeros(shape[0])
        z = np.arange(shape[2]) if meshtype == 'rectangular' else np.zeros(shape[0])
        generated = True
    elif x is None or y is None or z is None:
        raise ValueError("x, y, and z must either all be given or all omitted. ")
    return x, y, z, generated

def _generate_valueunits_list(valueunits, valuedim):
    if len(valueunits) == 0:
        valueunits = ["1" for a in range(valuedim)]
    elif len(valueunits) == 1:
        valueunits = [valueunits[0] for a in range(valuedim)]
    elif len(valueunits) != valuedim:
        raise ValueError("Length of valueunits must be 1 or match"
                         " the number of data dimensions. ")
    return shlex.join(valueunits)

def _generate_valuelabels_list(valuelabels, valuedim):
    if len(valuelabels) == 0:
        valuelabels = [f"value_{n}" for n in range(valuedim)]
    elif len(valuelabels) != valuedim:
        raise ValueError("Length of valuelabels must match "
                         "the number of data dimensions. ")
    return shlex.join(valuelabels)

def _shape_desc(desc):
    s = "# desc: OVF file generated by ovf2io.py."
    for line in desc:
        s += "\n# desc: " + line
    return s

def _make_header(header, representation):
    rep = {"text": "text", "bin4": "Binary 4", "bin8": "Binary 8"}[representation]
    if header['meshtype'] == 'rectangular':
        frontmatter = _templates.rectangular_template
    else:
        frontmatter = _templates.irregular_template
    for key in header.keys():
        frontmatter = frontmatter.replace(f"[{key}]", str(header[key]))
    frontmatter = frontmatter.replace("[repr]", rep)
    return frontmatter

def _format_data(data, meshtype, x, y, z):
    if meshtype == "rectangular":
        reshaped = data.reshape((-1, data.shape[-1]), order='F')
    else:
        reshaped = np.zeros((data.shape[0], data.shape[1] + 3))
        reshaped[:, 0] = x
        reshaped[:, 1] = y
        reshaped[:, 2] = z
        reshaped[:, 3:] = data
    return reshaped

def _write_file(fname, frontmatter, representation, reshaped):
    with open(fname, "wb") as f:
        f.write(frontmatter.encode("utf-8"))
        binrep = {"bin4": ("<f", 1234567.0), "bin8": ("<d", 123456789012345.0)}
        if representation in binrep:
            f.write(struct.pack(*binrep[representation]))
            f.write(reshaped.astype(binrep[representation][0]).tobytes())
            f.write("\n".encode("utf-8"))
        else:
            np.savetxt(f, reshaped)
        rep = {"text": "text", "bin4": "Binary 4", "bin8": "Binary 8"}[representation]
        f.write(f"# End: Data {rep}".encode("utf-8"))
        f.write("\n# End: Segment".encode("utf-8"))

