from . import _utils as ut
import numpy as np
from pathlib import Path

def read_ovf(fname):
    """Returns a dictionary containing the information read from an .ovf file.
    
    The returned dictionary has three items: 

    1. 'data', the actual values. 
    2. 'coords', generated coordinates based on xmin, xstepsize, and xnodes etc. 
    3. 'metadata', all of the information contained in the header. 

    'data' contains an entry for each axis of data stored, for example 
    read_ovf(fname)['data']['Zeeman energy density'] or 
    read_ovf(fname)['data']['m_x']

    'coords' contains coordinates generated from the header info, for example 
    read_ovf(fname)['coords']['x']

    If meshtype is 'rectangular', the shape of each 'data' entry and each 'coords' 
    entry will be (xnodes, ynodes, znodes). If meshtype is 'irregular', each entry 
    will be 1-dimensional. 

    **Parameters**

    * **fname** : _str or Path_ <br />
    The filename. <br />

    **Returns**

    * **data** : _dict_ <br />
    A dictionary containing the data, metadata, and generated coordinates.

    """
    fname = Path(fname)
    with open(fname, "rb") as f:
        if not b"2.0" in next(f):
            raise ValueError("This file does not appear to be OVF 2.0. "
                             "ovf2io does not support older OVF formats. ")
        # Skip ahead to the header
        while b"# begin: header" not in next(f).lower():
            pass
        header = ut._parse_header(f)
        nbytes = ut._advance_to_data_block(f)
        data = ut._parse_data(f, header, nbytes)
    coords = ut._gen_coords(data, header)
    header['repr'] = "text" if nbytes is None else f"Binary {nbytes}"
    out = {
            'data': data,
            'coords': coords,
            'metadata': header
        }
    return out

def write_ovf(data, fname, x=None, y=None, z=None,
        title="title", desc=[], meshunit="m", 
        valueunits=[], valuelabels=[], representation="bin8"
    ):
    """Write data to an OOMMF Vector Field (.ovf) file. 

    **Parameters**

    * **data** : _ndarray_ <br />
    The values to be stored. If `meshtype = "rectangular"`, data must have shape
    (len(x), len(y), len(z), num_data_components); for example, (10, 10, 10, 3) if there
    are 10 x-samples, 10 y-samples, 10 z-samples, and 3 vector components. 
    If `meshtype = "rectangular"`, data must have shape (num_locations, num_data_components).

    * **fname** : _str or Path_ <br />
    The name of the file to write. Will be overwritten if it exists already. 
    Intermediate directories are not created automatically. 

    * **x** : _ndarray, optional_ <br />
    X-coordinates, 1-dimensional array. If not given, indices will be assigned.

    * **y** : _ndarray, optional_ <br />
    Y-coordinates, 1-dimensional array. If not given, indices will be assigned.

    * **z** : _ndarray, optional_ <br />
    Z-coordinates, 1-dimensional array. If not given, indices will be assigned.

    * **title** : _str, optional_ <br />
    Title to be stored in metadata. <br />
    Default is `title = "title"`.

    * **desc** : _list, optional_ <br />
    Multi-line description; each line should be a separate entry in the list. 

    * **meshunit** : _str, optional_ <br />
    The unit of x, y, and z coordinates. <br />
    Default is `meshunit = "m"`.

    * **valueunits** : _list, optional_ <br />
    The units for each value stored. Length should match the number of data components, 
    or can be length 1 if all data components share the same units. If not given, 
    all data components will be given units "1". 

    * **valuelabels** : _list, optional_ <br />
    Labels for each data component. Length should match the number of data components. 
    If not given, components will be labelled "value_0", "value_1", etc.

    * **representation** : _str, optional_ <br />
    The storage mode for the data itself. One of "text", "bin4", and "bin8". 
    Comments are allowed in text mode, which uses the UTF-8 encoding. 
    Note that the original specification specifies ASCII. <br />
    Default is `representation = "bin8"`. 
    """
    data = np.array(data)

    # Add a line to desc saying generated by ovf2io
    desc = ut._shape_desc(desc)
    valuedim = data.shape[-1]

    if len(data.shape) == 4:
        meshtype = "rectangular"
    elif len(data.shape) == 2:
        meshtype = "irregular"
    else:
        raise Exception("Data is not an expected shape. "
                        "Data should have shape (len(x), len(y), len(z), num_data_components) "
                        "for rectangular meshes and shape (num_points, num_data_components) "
                        "for irregular meshes. ")
    # Check meshtype is ok
    # if not meshtype in {"rectangular", "irregular"}:
        # raise ValueError("Meshtype must be either 'rectangular' or 'irregular'.")

    # Generate coordinates and unit
    x, y, z, generated = ut._generate_coordinates(x, y, z, data.shape, meshtype)
    if generated: meshunit = "pt"

    # Generate valueunits and valuelabels
    valueunits = ut._generate_valueunits_list(valueunits, valuedim)
    valuelabels = ut._generate_valuelabels_list(valuelabels, valuedim)

    header = {
            "title": title, "desc": desc, "meshunit": meshunit, "meshtype": meshtype,
            "valueunits": valueunits, "valuelabels": valuelabels, "valuedim": valuedim,
        }
    if meshtype == "rectangular":
        header.update({
            "xbase": x[0], "ybase": y[0], "zbase": z[0], 
            "xstepsize": abs(x[1]-x[0]), "ystepsize": abs(y[1]-y[0]), "zstepsize": abs(z[1]-z[0]),
            "xnodes": len(x), "ynodes": len(y), "znodes": len(z),
            "xmin": np.min(x) - 0.5 * abs(x[1]-x[0]), "xmax": np.max(x) + 0.5 * abs(x[1]-x[0]),
            "ymin": np.min(y) - 0.5 * abs(y[1]-y[0]), "ymax": np.max(y) + 0.5 * abs(y[1]-y[0]),
            "zmin": np.min(z) - 0.5 * abs(z[1]-z[0]), "zmax": np.max(z) + 0.5 * abs(z[1]-z[0]),
        })
    elif meshtype == "irregular":
        header.update({
            "pointcount": data.shape[0],
            "xmin": np.min(x), "xmax": np.max(x),
            "ymin": np.min(y), "ymax": np.max(y),
            "zmin": np.min(z), "zmax": np.max(z),
        })
    if not representation.lower() in {"text", "bin4", "bin8"}:
        raise ValueError("Representation must be either 'text', 'bin4', or 'bin8'.")
    reshaped = ut._format_data(data, meshtype, x, y, z)
    frontmatter = ut._make_header(header, representation)
    ut._write_file(fname, frontmatter, representation, reshaped)
