#!/usr/bin/env python3
import numpy as np
import yaml
import h5py
import os
import warnings

warnings.simplefilter("always", DeprecationWarning)
warnings.formatwarning = lambda msg, *args, **kwargs: f"{msg}\n"

PARSER_VERSION = "0.3.0"

SUPPORTED_REPRESENTATIONS = {"canonical"}  # extend this set when adding new representations
SUPPORTED_CONTENT = {"all", "metadata", "fields", "wavefunctions"} # New supported content types

#----------------------------------------------------------
def _ensure_version(metadata):
    """Ensure the metadata contains the current parser version."""
    metadata["version"] = PARSER_VERSION
    return metadata


def _get_representation(wf_meta):
    rep = wf_meta.get("representation", "canonical")

    if not isinstance(rep, str):
        raise ValueError("representation must be a string")

    rep_norm = rep.lower().strip()

    if rep_norm not in SUPPORTED_REPRESENTATIONS:
        raise ValueError(f"Unknown wavefunction representation '{rep}'. "
                         f"Supported: {SUPPORTED_REPRESENTATIONS}")

    return rep_norm

#----------------------------------------------------------
def _sanitize_metadata(data):
    """
    Recursively convert NumPy objects (arrays, scalars) in a dictionary
    to native Python types (list, float, int) for clean YAML dumping.
    """
    if isinstance(data, dict):
        return {k: _sanitize_metadata(v) for k, v in data.items()}
    elif isinstance(data, (list, tuple)):
        # Handle lists/tuples, which often contain the field metadata lists
        return [_sanitize_metadata(item) for item in data]
    elif isinstance(data, np.bytes_):
        # Convert np.bytes_ to string
        return data.decode()
    elif isinstance(data, np.ndarray):
        # Convert NumPy arrays to lists of native Python types
        return data.tolist()
    elif isinstance(data, np.generic):
        # Convert NumPy scalar types (e.g., np.float64) to native Python types (float)
        return data.item()
    else:
        return data

#----------------------------------------------------------
def ixleft_to_ixright(array):
    ''' 
    Takes as input a numpy array as input with the index order 
    [ix,iy,iz,...] and any internal storage order.
    Returns a numpy array ready for hdf5/binary storage with the indices 
    [...,iz,iy,ix] and row-major ordering (ix varies the faster).
    '''
    result = np.array(0)

    a = np.asarray(array)
    if a.flags['C_CONTIGUOUS']:
        result = np.ascontiguousarray(np.transpose(a))
    elif a.flags['F_CONTIGUOUS']:
        reversed_shape = a.shape[::-1]
        serialized = np.ravel(a, order='A')
        result = np.reshape(serialized, reversed_shape)
    else:
        raise ValueError("Unkown internal storage for the array")

    return result

#----------------------------------------------------------
def ixright_to_ixleft(array):
    ''' 
    Takes as input an array read from an hdf5 file with the index order
    [...,iz,iy,ix] and a C_CONTIGUOUS internal storage.
    Returns a numpy array with the index order [ix,iy,iz,...] 
    and a F_CONTIGUOUS internal storage order (ix varies the faster).
    Only the indices order is reversed but the internal data is left unchanged.
    '''
    a = np.asarray(array)
    if not a.flags['C_CONTIGUOUS']:
        raise ValueError("Internal storage order not yet supported")

    reversed_shape = a.shape[::-1]
    serialized = np.ravel(a, order='A')
    return np.reshape(serialized, reversed_shape, order='F')

#----------------------------------------------------------
def expand_field_dims(array):
    '''
    Add the frame and number of components dimensions to 
    a field array if needed. The output array have
    indices of the form [ix,iy,iz,ic,it] where
    ic stands for the component index and it for the frame (time) index.

    Example: input array.shape  = (nx,ny,nz)
             output array.shape = (nx,ny,nz,1,1)
    '''
    res = array
    n_dims = len(array.shape)
    if n_dims == 3:
        res = np.expand_dims(array, (3,4))
    elif n_dims == 4:
        res = np.expand_dims(array, 4)
    elif n_dims == 5:
        res = array
    else:
        raise ValueError("Wrong number of axes in the input array")
    return res 

#----------------------------------------------------------
def expand_uv_dims(array):
    '''
    Add the frame number dimension to a u or v array if needed.
    The output array have a shape (nw, nt) where
    nw stands for the number of single-particle states
    and nt for the number of frames (times).

    Example: input array.shape  = (nw)
             output array.shape = (nw,nt)
    '''
    res = array
    n_dims = len(array.shape)
    if n_dims == 1:
        res = np.expand_dims(array, 1)
    else:
        raise ValueError("Wrong number of axes in the input array")
    return res 

#----------------------------------------------------------
def expand_states_dims(array):
    '''
    Add the frame number dimension to a states array array if needed.
    The output array have a shape (nx,ny,nz,ns,nw,nt) where
    nw stands for the number of single-particle states
    and nt for the number of frames (times).

    Example: input array.shape  = (nx,ny,nz,ns,nw)
             output array.shape = (nx,ny,nz,ns,nw,nt)
    '''
    res = array
    n_dims = len(array.shape)
    if n_dims == 5:
        res = np.expand_dims(array, 5)
    else:
        raise ValueError("Wrong number of axes in the input array")
    return res 

#----------------------------------------------------------
def reduce_last_dims(array, n_dims=1):
    '''
    @param array  input array
    @param n_dims number of dimensions to be tentatively removed

    Remove the right-most dimensions of an array if they contain only
    1 element.

    Example: n_dims = 2
             input array.shape  = (nx,ny,nz,1,1)
             output array.shape = (nx,ny,nz)

             n_dims = 2
             input array.shape  = (nx,ny,nz,2,1)
             output array.shape = (nx,ny,nz,2)

             n_dims = 2
             input array.shape  = (nx,ny,nz,1,2)
             output array.shape = (nx,ny,nz,1,2)
    '''
    shape = array.shape
    new_shape = shape
    for i in range(n_dims):
        if new_shape[-1] == 1:
            new_shape = new_shape[:-1]

    return np.reshape(array, new_shape, order='A')

#----------------------------------------------------------
def write_hdf5_root(data, group):
    '''
    Write the root group
    '''
    metadata = data['metadata']

    for key in ['N', 'Z', 'description', 'version']:
        if key in metadata:
            group.attrs[key] = metadata[key]

#----------------------------------------------------------
def write_hdf5_discretization(data, group):
    '''
    Write the discretization group
    '''
    metadata = data['metadata']

    # --- ensure t0 exists ---
    if 't0' not in metadata:
        metadata['t0'] = 0

    # --- ensure t0 exists ---
    if 'frame' not in metadata:
        metadata['frame'] = 1

    # deprecated key names
    group.attrs['tmin'] = metadata['t0']
    group.attrs['nt']   = metadata['frame']

    # completion
    if not 'xmin' in metadata:
        metadata['xmin'] = -0.5 * metadata['nx'] * metadata['dx']
    if not 'ymin' in metadata:
        metadata['ymin'] = -0.5 * metadata['ny'] * metadata['dy']
    if not 'zmin' in metadata:
        metadata['zmin'] = -0.5 * metadata['nz'] * metadata['dz']

    # next gen keynames
    attribute_keys = ['dx', 'dy', 'dz', 'dt', 
                      'nt', 'nx', 'ny', 'nz',
                      'xmin', 'ymin', 'zmin', 'tmin',
                      'mesh_type', 'symmetry']

    for key in attribute_keys:
        if key in metadata:
            group.attrs[key] = metadata[key]

#----------------------------------------------------------
def write_hdf5_fields(data, group):
    '''
    Write the fields group
    '''
    metadata = data['metadata']
    if not 'field' in metadata:
        return

    fields = metadata['field']
    for fld in fields:
        key = fld['name']
        if key not in data:
            continue

        # Field array
        array = data[key]
        array = expand_field_dims(array)
        array = ixleft_to_ixright(array)
        dset  = group.create_dataset(key, data=array, compression='gzip')

        # Field attributes
        sanitized_fld = _sanitize_metadata(fld.copy())
        if 'unit' in sanitized_fld:
            dset.attrs['unit'] = sanitized_fld['unit']
        if 'ev8_sym' in sanitized_fld:
            dset.attrs['ev8_reflexion_signs'] = sanitized_fld['ev8_sym']

#----------------------------------------------------------
def write_hdf5_wavefunction(data, group):
    '''
    Write wavefunction group
    '''
    metadata = data['metadata']
    if not 'wavefunction' in metadata:
        return

    wf_meta = metadata['wavefunction']
    rep = _get_representation(wf_meta)
    group.attrs['representation'] = rep

    match rep:
        case "canonical":
            # --- Neutron block ---
            n_group = group.create_group('neutron')
            
            if 'symblocks_n' in wf_meta:
                n_group.attrs['symblocks'] = wf_meta['symblocks_n']
    
            if 'statesn' in data:
                array = expand_states_dims(data['statesn'])
                array = ixleft_to_ixright(array)
                n_group.create_dataset('states', data=array, compression='gzip')
            if 'un' in data:
                array = expand_uv_dims(data['un'])
                array = ixleft_to_ixright(array)
                n_group.create_dataset('u', data=array, compression='gzip')
            if 'vn' in data:
                array = expand_uv_dims(data['vn'])
                array = ixleft_to_ixright(array)
                n_group.create_dataset('v', data=array, compression='gzip')

            # --- Proton block ---
            p_group = group.create_group('proton')

            if 'symblocks_p' in wf_meta:
                p_group.attrs['symblocks'] = wf_meta['symblocks_p']
            
            if 'statesp' in data:
                array = expand_states_dims(data['statesp'])
                array = ixleft_to_ixright(array)
                p_group.create_dataset('states', data=array, compression='gzip')
            if 'up' in data:
                array = expand_uv_dims(data['up'])
                array = ixleft_to_ixright(array)
                p_group.create_dataset('u', data=array, compression='gzip')
            if 'vp' in data:
                array = expand_uv_dims(data['vp'])
                array = ixleft_to_ixright(array)
                p_group.create_dataset('v', data=array, compression='gzip')

        case _:
            raise ValueError(f"Unsupported representation '{rep}'.")

#----------------------------------------------------------
def write_hdf5(data, file_name):
    """
    Write a Bogoliubov state and associated fields to an HDF5 file.
    """
    metadata = data['metadata']
    metadata = _ensure_version(metadata)

    odir = os.path.dirname(file_name) or '.'
    os.makedirs(odir, exist_ok=True)

    with h5py.File(file_name, 'w') as h5f:

        write_hdf5_root(data, h5f)

        discretization = h5f.create_group('discretization')
        write_hdf5_discretization(data, discretization)

        if 'field' in metadata:
            fields = h5f.create_group('fields')
            write_hdf5_fields(data, fields)

        if 'wavefunction' in metadata:
            wavefunction = h5f.create_group('wavefunction')
            write_hdf5_wavefunction(data, wavefunction)

#----------------------------------------------------------
def read_hdf5_root(group, data, content='all'):
    '''
    Read the root group
    '''
    metadata = data['metadata']
    for key in ['N', 'Z', 'description', 'version']:
        if key in group.attrs:
            metadata[key] = group.attrs[key]

#----------------------------------------------------------
def read_hdf5_fields(group, data, content='all'):
    '''
    Read the fields group
    '''
    if not content in ['all', 'fields']:
        return

    metadata_list = []
    for key, dset in group.items():
        # Read data
        array = ixright_to_ixleft(dset)
        array = reduce_last_dims(array, n_dims=2)
        data[key] = array

        # Store metadata
        fld_meta = {}
        for k, v in dset.attrs.items():
            if isinstance(v, np.generic):
                v = v.item()
            fld_meta[k] = v
        fld_meta['name']   = key
        fld_meta['suffix'] = f"_{key}.bin"
        metadata_list.append(fld_meta)
    data['metadata']['field'] = metadata_list

#----------------------------------------------------------
def read_hdf5_wavefunction(group, data, content='all'):
    '''
    Read the wavefunction group
    '''
    # representation from file (optional)

    # rep_in_file = group.attrs.get("representation", None)
    # if 'wavefunction' in metadata:
    #     # prefer metadata
    #     rep = _get_representation(metadata['wavefunction'])
    #     if rep_in_file and rep_in_file != rep:
    #         warnings.warn(
    #             f"Representation mismatch: metadata='{rep}', file='{rep_in_file}'. "
    #             f"Using metadata.",
    #             DeprecationWarning
    #         )
    # else:
    #     # fallback if metadata lacks wf section
    #     rep = rep_in_file or "canonical"
    #     metadata['wavefunction'] = {'representation': rep} # Add minimal wf metadata

    # Read wavefunction metadata
    if not content in ['all', 'wavefunctions', 'metadata']:
        return

    wf_metadata = {}
    wf_metadata['representation'] = group.attrs.get("representation", None)
    wf_metadata['suffix'] = '_state'
    wf_metadata['n_neutron_states'] = group['neutron/states'].shape[1]
    wf_metadata['n_proton_states'] = group['proton/states'].shape[1]
    
    # --- read symblocks from attributes ---
    if 'neutron' in group:
        n_group = group['neutron']
        if 'symblocks' in n_group.attrs:
            wf_metadata['symblocks_n'] = list(n_group.attrs['symblocks'])
    if 'proton' in group:
        p_group = group['proton']
        if 'symblocks' in p_group.attrs:
            wf_metadata['symblocks_p'] = list(p_group.attrs['symblocks'])
    
    data['metadata']['wavefunction'] = _sanitize_metadata(wf_metadata)

    # Read wavefunction data
    if not content in ['all', 'wavefunctions']:
        return

    rep = data['metadata']['wavefunction']['representation']
    match rep:
        case "canonical":
            if 'neutron' in group:
                n_group = group['neutron']
                if 'states' in n_group:
                    array = ixright_to_ixleft(n_group['states'])
                    array = reduce_last_dims(array, n_dims=1)
                    data['statesn'] = array
                if 'u' in n_group:
                    array = ixright_to_ixleft(n_group['u'])
                    array = reduce_last_dims(array, n_dims=1)
                    data['un'] = array
                if 'v' in n_group:
                    array = ixright_to_ixleft(n_group['v'])
                    array = reduce_last_dims(array, n_dims=1)
                    data['vn'] = array

            if 'proton' in group:
                p_group = group['proton']
                if 'states' in p_group:
                    array = ixright_to_ixleft(p_group['states'])
                    array = reduce_last_dims(array, n_dims=1)
                    data['statesp'] = array
                if 'u' in p_group:
                    array = ixright_to_ixleft(p_group['u'])
                    array = reduce_last_dims(array, n_dims=1)
                    data['up'] = array
                if 'v' in p_group:
                    array = ixright_to_ixleft(p_group['v'])
                    array = reduce_last_dims(array, n_dims=1)
                    data['vp'] = array

        case _:
            raise ValueError(f"Unsupported representation '{rep}'.")

#----------------------------------------------------------
def read_hdf5_discretization(group, data, content='all'):
    '''
    Read the wavefunction group
    '''
    if not content in ['all', 'metadata', 'wavefunctions']:
        return

    metadata = data['metadata']
    for key in ['nx', 'ny', 'nz', 'nt',
                'dx', 'dy', 'dz', 'dt',
                'xmin', 'ymin', 'zmin', 'tmin',
                'mesh_type', 'symmetry']:
        if key in group.attrs:
            metadata[key] = group.attrs[key]

    # Deprecated keys
    metadata['frame'] = metadata['nt']
    metadata['t0']    = metadata['tmin']

#----------------------------------------------------------
def read_hdf5(filepath, content="all"):
    """
    Read a Bogoliubov state and fields from an HDF5 file.
    Restores a metadata dictionary compatible with the binary/YAML format.

    content: "all" (default), "metadata", "fields", or "wavefunctions"
    """
    if content not in SUPPORTED_CONTENT:
        raise ValueError(f"Unknown content type '{content}'. Supported: {SUPPORTED_CONTENT}")

    data = {}
    if not 'metadata' in data:
        data['metadata'] = {}
    metadata = data['metadata']

    prefix = os.path.splitext(os.path.basename(filepath))[0]
    metadata['prefix'] = prefix

    with h5py.File(filepath, 'r') as h5f:

        read_hdf5_root(h5f, data, content)

        if 'discretization' in h5f:
            read_hdf5_discretization(h5f['discretization'], data, content)

        if 'fields' in h5f:
            read_hdf5_fields(h5f['fields'], data, content)

        if 'wavefunction' in h5f:
            read_hdf5_wavefunction(h5f['wavefunction'], data, content)

    data['metadata'] = _sanitize_metadata(metadata.copy())

    return data

#----------------------------------------------------------
def write_yaml(data, filepath):
    """
    Write Bogoliubov state and fields to YAML + binary format.
    Ensures binary arrays are written in strict Fortran (column-major) order.
    """
    metadata = _ensure_version(data['metadata'])
    prefix = metadata['prefix']
    odir = os.path.dirname(filepath) or '.'
    os.makedirs(odir, exist_ok=True)

    # Sanitize metadata before dumping to prevent NumPy-specific serialization
    sanitized_metadata = _sanitize_metadata(metadata.copy())

    # Write YAML metadata
    with open(filepath, 'w') as file:
        yaml.dump(sanitized_metadata, file)

    # Write fields
    if 'field' in metadata:
        fields = metadata['field']
        for fld in fields:
            key = fld['name']
            if key not in data:
                continue

            array = ixleft_to_ixright(data[key])
            f_name = os.path.join(odir, prefix + fld['suffix'])
            with open(f_name, 'wb') as f:
                # f.write(np.asfortranarray(array).tobytes(order='F'))
                f.write(array.tobytes())

    #------------------------------------------------------
    # Wavefunctions
    if 'wavefunction' in metadata:
        wf_meta = metadata['wavefunction']
        rep = _get_representation(wf_meta)

        f_name = os.path.join(odir, prefix + wf_meta['suffix'])

        match rep:
            case "canonical":
                with open(f_name, 'wb') as file:
                    # Only write if all required keys are present
                    required_keys = ['statesn', 'statesp', 'un', 'vn', 'up', 'vp']
                    if all(k in data for k in required_keys):
                        buffer = (
                            data['statesn'].tobytes(order='F') +
                            data['statesp'].tobytes(order='F') +
                            data['un'].tobytes() +
                            data['vn'].tobytes() +
                            data['up'].tobytes() +
                            data['vp'].tobytes()
                        )
                        file.write(buffer)
                    else:
                        print(f"Warning: Skipping wavefunction binary write for '{f_name}' due to missing data.")

            case _:
                raise ValueError(f"Unsupported wavefunction representation '{rep}'.")


#----------------------------------------------------------
def read_yaml(odir, header, content="all"):
    """
    Read a Bogoliubov state and fields (potentials, densities)
    from YAML + binary format.

    content: "all" (default), "metadata", "fields", or "wavefunctions"
    """
    if content not in SUPPORTED_CONTENT:
        raise ValueError(f"Unknown content type '{content}'. Supported: {SUPPORTED_CONTENT}")

    data = {}
    with open(header, 'r') as file:
        metadata = yaml.safe_load(file)
        metadata = _ensure_version(metadata)
        data['metadata'] = metadata

    if 'variables' in metadata:
        warnings.warn("YAML key 'variables' is deprecated; using 'field' instead.", DeprecationWarning)
        metadata['field'] = metadata.pop('variables')

    if content == 'metadata':
        return data

    prefix = metadata['prefix']
    nx, ny, nz = metadata['nx'], metadata['ny'], metadata['nz']
    symmetry = metadata.get('symmetry', '')
    if symmetry == 'ev8':
        nx, ny, nz = nx // 2, ny // 2, nz // 2

    # Read Fields
    if content in ("all", "fields"):
        n_frames = metadata.get('frame', 1)
        if 'field' in metadata:
            fields = metadata['field']
            valid_fields = []
            for fld in fields:
                key = fld['name']
                f_name = os.path.join(odir, prefix + fld['suffix'])
                nc = fld['n_components']
                try:
                    with open(f_name, 'rb'):
                        pass
                except FileNotFoundError:
                    print(f"Warning: missing field file '{f_name}', skipping.")
                    continue

                if n_frames == 1:
                    shape = (nx, ny, nz, nc)
                    if nc == 1:
                        shape = (nx, ny, nz)
                else:
                    shape = (nx, ny, nz, nc, n_frames)
                    if nc == 1:
                        shape = (nx, ny, nz, n_frames)

                data[key] = np.reshape(np.fromfile(f_name, dtype=np.float64), shape, order='F')
                valid_fields.append(fld)

            metadata['field'] = valid_fields

    # Read Wavefunctions
    if content in ("all", "wavefunctions"):
        if 'wavefunction' in metadata:
            wf_meta = metadata['wavefunction']
            rep = _get_representation(wf_meta)

            f_name = os.path.join(odir, prefix + wf_meta['suffix'])
            nw_n = wf_meta['n_neutron_states']
            nw_p = wf_meta['n_proton_states']

            match rep:
                case "canonical":
                    try:
                        with open(f_name, 'rb') as file:

                            statesn_bsize = nw_n * nx * ny * nz * 2 * np.complex128(0).nbytes
                            buffer = file.read(statesn_bsize)
                            data['statesn'] = np.reshape(np.frombuffer(buffer, dtype=np.complex128),
                                                         (nx, ny, nz, 2, nw_n), order='F')

                            statesp_bsize = nw_p * nx * ny * nz * 2 * np.complex128(0).nbytes
                            buffer = file.read(statesp_bsize)
                            data['statesp'] = np.reshape(np.frombuffer(buffer, dtype=np.complex128),
                                                         (nx, ny, nz, 2, nw_p), order='F')

                            un_bsize = nw_n * np.complex128(0).nbytes
                            data['un'] = np.frombuffer(file.read(un_bsize), dtype=np.complex128)
                            data['vn'] = np.frombuffer(file.read(un_bsize), dtype=np.complex128)

                            up_bsize = nw_p * np.complex128(0).nbytes
                            data['up'] = np.frombuffer(file.read(up_bsize), dtype=np.complex128)
                            data['vp'] = np.frombuffer(file.read(up_bsize), dtype=np.complex128)

                    except FileNotFoundError:
                        print(f"Warning: missing wavefunction file '{f_name}', skipping wavefunction read.")
                        # If file is missing, we still return the metadata
                        pass
                case _:
                    raise ValueError(f"Unsupported representation '{rep}'.")

    return data


#----------------------------------------------------------
# Unified interface
def read(filepath, content="all"):
    """
    Auto-detect file type and read accordingly.
      - .yaml/.yml -> YAML format
      - .h5/.hdf5  -> HDF5 format

    content: "all" (default), "metadata", "fields", or "wavefunctions"
    """
    if content not in SUPPORTED_CONTENT:
        raise ValueError(f"Unknown content type '{content}'. Supported: {SUPPORTED_CONTENT}")
        
    if filepath.endswith(('.yaml', '.yml')):
        odir = os.path.dirname(filepath) or '.'
        return read_yaml(odir, filepath, content=content)
    elif filepath.endswith(('.h5', '.hdf5')):
        return read_hdf5(filepath, content=content)
    else:
        raise ValueError(f"Unrecognized file extension for '{filepath}'.")


def write(data, filepath):
    """
    Unified writer. File format is determined by the extension of `filepath`.
    """
    ext = os.path.splitext(filepath)[1].lower()
    data['metadata'] = _ensure_version(data['metadata'])
    if ext in ('.yaml', '.yml'):
        write_yaml(data, filepath)
    elif ext in ('.h5', '.hdf5'):
        write_hdf5(data, filepath)
    else:
        raise ValueError(f"Unknown file extension for output: '{filepath}'. Expected .yaml/.yml or .h5/.hdf5.")
