import os
import h5py
import numpy as np
import pandas as pd
import pynibs


def create_stimsite_hdf5(fn_exp, fn_hdf, conditions_selected=None, sep="_", merge_sites=False, fix_angles=False,
                         data_dict=None, conditions_ignored=None):
    """
    Reads results_conditions and creates an hdf5/xdmf pair with condition-wise centers of stimulation sites and
    coil directions as data.

    Parameters
    ----------
    fn_exp : str
        Path to results.csv.
    fn_hdf : str
        Path where to write file. Gets overridden if already existing.
    conditions_selected : str or list of str, optional
        List of conditions returned by the function, the others are omitted.
        If None, all conditions are returned.
    sep: str, default: "_"
        Separator between condition label and angle (e.g. M1_0, or M1-0).
    merge_sites : bool
        If true, only one coil center per site is generated.
    fix_angles : bool
        rename 22.5 -> 0, 0 -> -45, 67.5 -> 90, 90 -> 135.
    data_dict : dict ofnp.ndarray of float [n_stimsites] (optional), default: None
        Dictionary containing data corresponding to the stimulation sites (keys).
    conditions_ignored : str or list of str, optional
        Conditions, which are not going to be included in the plot.

    Returns
    -------
    <Files> : hdf5/xdmf file pair
        Contains information about condition-wise stimulation sites and coil directions (fn_hdf)

    Example
    -------
    .. code-block:: python

        pynibs.create_stimsite_hdf5('/exp/1/experiment_corrected.csv',
                                    '/stimsite', True, True)
    """
    assert not fn_hdf.endswith('/')

    exp = pynibs.expio.read_csv(fn_exp)

    exp_cond = pynibs.expio.sort_by_condition(exp, conditions_selected=conditions_selected)  # []

    # get the unique conditions in the correct order
    conds = [c['condition'][0] for c in exp_cond]

    # remove conds
    conds_temp = []
    exp_cond_temp = []

    if type(conditions_ignored) is not list:
        conditions_ignored = [conditions_ignored]

    for i_c, c in enumerate(conds):
        ignore = False
        for ci in list(conditions_ignored):
            if c == ci:
                ignore = True

        if not ignore:
            conds_temp.append(conds[i_c])
            exp_cond_temp.append(exp_cond[i_c])

    exp_cond = exp_cond_temp
    conds = conds_temp

    # hardcoded row #3 is condition
    cond_idx = np.linspace(0, len(exp_cond), 1)[:, np.newaxis]

    centers = []
    m0 = []
    m1 = []
    m2 = []

    for i_cond in range(len(exp_cond)):
        centers.append(exp_cond[i_cond]['coil_mean_matrix'][0][0:3, 3])
        m0.append(exp_cond[i_cond]['coil_mean_matrix'][0][0:3, 0])
        m1.append(exp_cond[i_cond]['coil_mean_matrix'][0][0:3, 1])
        m2.append(exp_cond[i_cond]['coil_mean_matrix'][0][0:3, 2])

    # split conds to angles and sites: M1_90 -> M1, 90
    angles = np.array([sp.split(sep)[-1] for sp in conds]).astype(np.float64)
    sites = np.array([sp.split(sep)[0] for sp in conds])
    sites_unique = np.unique(sites)

    # average the center positions of a stimulation site over all orientations
    if merge_sites:

        # generate sites dict
        centers_sites = dict()

        for site in sites_unique:
            centers_sites[site] = []

        # gather all orientations and put them to the corresponding sites
        for i_cond, site in enumerate(sites):
            centers_sites[site].append(exp_cond[i_cond]['coil_mean_matrix'][0][0:3, 3])

        # determine average position over all orientations for each site
        for site in sites_unique:
            centers_sites[site] = np.mean(np.vstack(centers_sites[site]), axis=0)

        # write it back to centers
        for i_cond, site in enumerate(sites):
            centers[i_cond] = centers_sites[site]

    centers = np.vstack(centers)
    m0 = np.vstack(m0)
    m1 = np.vstack(m1)
    m2 = np.vstack(m2)

    # enumerate sites, as paraview does not plot string array data
    sites_idx = np.array(list(range(len(sites))))[:, np.newaxis]

    angles[angles[:] == 675.] = 67.5
    angles[angles[:] == 225.] = 22.5

    if fix_angles:
        # rename wrong angle names
        angles_cor = np.copy(angles)
        angles_cor[angles == 0] = -45.
        angles_cor[angles == 22.5] = 0.
        angles_cor[angles == 67.5] = 90.
        angles_cor[angles == 90] = 135.
        angles = angles_cor

    # write hdf5 file
    if not fn_hdf.endswith('.hdf5'):
        fn_hdf += '.hdf5'
    f = h5py.File(fn_hdf, 'w')
    f.create_dataset('centers', data=centers.astype(np.float64))
    f.create_dataset('m0', data=m0.astype(np.float64))
    f.create_dataset('m1', data=m1.astype(np.float64))
    f.create_dataset('m2', data=m2.astype(np.float64))
    f.create_dataset('cond', data=np.string_(conds))  # this is a string array, not xdmf compatible
    f.create_dataset('cond_idx', data=cond_idx)
    f.create_dataset('angles', data=angles)
    f.create_dataset('sites', data=np.string_(sites))  # this is a string array, not xdmf compatible
    f.create_dataset('sites_idx', data=sites_idx)

    data = None
    if data_dict is not None:
        data = np.zeros((len(list(data_dict.keys())), 1))
        for i_data, cond in enumerate(conds):
            data[i_data, 0] = data_dict[cond]
        f.create_dataset('data', data=data)
    f.close()

    # write .xdmf file
    f = open(fn_hdf[:-4] + 'xdmf', 'w')
    fn_hdf = os.path.basename(fn_hdf)  # relative links

    # header
    f.write('<?xml version="1.0"?>\n')
    f.write('<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>\n')
    f.write('<Xdmf Version="2.0" xmlns:xi="http://www.w3.org/2001/XInclude">\n')
    f.write('<Domain>\n')
    f.write('<Grid\nCollectionType="Spatial"\nGridType="Collection"\nName="Collection">\n')

    # one grid for coil dipole nodes...store data hdf5.
    #######################################################
    f.write('<Grid Name="stimsites" GridType="Uniform">\n')
    f.write('<Topology NumberOfElements="' + str(centers.shape[0]) +
            '" TopologyType="Polyvertex" Name="Tri">\n')
    f.write('<DataItem Format="XML" Dimensions="' + str(centers.shape[0]) + ' 1">\n')
    # f.write(hdf5_fn + ':' + path + '/triangle_number_list\n')
    np.savetxt(f, list(range(centers.shape[0])), fmt='%d', delimiter=' ')  # 1 2 3 4 ... N_Points
    f.write('</DataItem>\n')
    f.write('</Topology>\n')

    # nodes
    f.write('<Geometry GeometryType="XYZ">\n')
    f.write('<DataItem Format="HDF" Dimensions="' + str(centers.shape[0]) + ' 3">\n')
    f.write(fn_hdf + ':' + '/centers\n')
    f.write('</DataItem>\n')
    f.write('</Geometry>\n')

    # data
    # dipole magnitude
    # the 4 vectors
    for i in range(3):
        f.write('<Attribute Name="dir_' + str(i) + '" AttributeType="Vector" Center="Cell">\n')
        f.write('<DataItem Format="HDF" Dimensions="' + str(centers.shape[0]) + ' 3">\n')
        f.write(fn_hdf + ':' + '/m' + str(i) + '\n')
        f.write('</DataItem>\n')
        f.write('</Attribute>\n\n')

    #  angles
    f.write('<Attribute Name="angles" AttributeType="Scalar" Center="Cell">\n')
    f.write('<DataItem Format="HDF" Dimensions="' + str(centers.shape[0]) + ' 1">\n')
    f.write(fn_hdf + ':' + '/angles\n')
    f.write('</DataItem>\n')
    f.write('</Attribute>\n\n')

    #  data
    if data_dict is not None:
        f.write('<Attribute Name="data" AttributeType="Scalar" Center="Cell">\n')
        f.write('<DataItem Format="HDF" Dimensions="' + str(data.shape[0]) + ' 1">\n')
        f.write(fn_hdf + ':' + '/data\n')
        f.write('</DataItem>\n')
        f.write('</Attribute>\n\n')

    # site idx
    f.write('<Attribute Name="sites_idx" AttributeType="Scalar" Center="Cell">\n')
    f.write('<DataItem Format="HDF" Dimensions="' + str(centers.shape[0]) + ' 1">\n')
    f.write(fn_hdf + ':' + '/sites_idx\n')
    f.write('</DataItem>\n')
    f.write('</Attribute>\n\n')

    f.write('</Grid>\n')
    # end coil dipole data

    # footer
    f.write('</Grid>\n')
    f.write('</Domain>\n')
    f.write('</Xdmf>\n')
    f.close()


def create_stimsite_from_list(fn_hdf, poslist, datanames=None, data=None, overwrite=False):
    """
    This takes a list of matsimnibs-style coil position and orientations and creates an .hdf5 + .xdmf tuple
    for all positions.

    Centers and coil orientations are written to disk, with optional data for each coil configuration.

    Parameters
    ----------
    fn_hdf: str
        Filename for the .hdf5 file. The .xdmf is saved with the same basename.
        Folder should already exist.
    poslist: list of np.ndarray
        (4,4) Positions.
    datanames: str or list of str, optional
        Dataset names for ``data``.
    data: np.ndarray, optional
        Dataset array with shape = ``(len(poslist.pos), len(datanames())``.
    overwrite : bool, defaul: False
        Overwrite existing files.
    """
    centers = []
    m0 = []
    m1 = []
    m2 = []
    if data is not None:
        assert isinstance(data, np.ndarray)

    for lst in poslist:
        centers.append(lst[0:3, 3])
        m0.append(lst[0:3, 0])
        m1.append(lst[0:3, 1])
        m2.append(lst[0:3, 2])

    centers = np.vstack(centers)
    m0 = np.vstack(m0)
    m1 = np.vstack(m1)
    m2 = np.vstack(m2)

    write_coil_pos_hdf5(fn_hdf, centers, m0, m1, m2, datanames=datanames, data=data, overwrite=overwrite)


def create_stimsite_from_tmslist(fn_hdf, poslist, datanames=None, data=None, overwrite=False):
    """
    This takes a :py:class:simnibs.sim_struct.TMSLIST from simnibs and creates an .hdf5 + .xdmf tuple for all positions.

    Centers and coil orientations are written to disk, with optional data for each coil configuration.

    Parameters
    ----------
    fn_hdf: str
        Filename for the .hdf5 file. The .xdmf is saved with the same basename.
        Folder should already exist.
    poslist: simnibs.sim_struct.TMSLIST
        poslist.pos[*].matsimnibs have to be set.
    datanames: str or list of str, optional
        Dataset names for ``data``.
    data: np.ndarray, optional
        Dataset array with shape = ``(len(poslist.pos), len(datanames())``.
    overwrite : bool, default: False
        Overwrite existing files
    """
    centers = []
    m0 = []
    m1 = []
    m2 = []
    assert poslist.pos
    if data is not None:
        assert isinstance(data, np.ndarray)
    for pos in poslist.pos:
        assert pos.matsimnibs is not None
        pos.matsimnibs = np.array(pos.matsimnibs)
        centers.append(pos.matsimnibs[0:3, 3])
        m0.append(pos.matsimnibs[0:3, 0])
        m1.append(pos.matsimnibs[0:3, 1])
        m2.append(pos.matsimnibs[0:3, 2])

    centers = np.vstack(centers)
    m0 = np.vstack(m0)
    m1 = np.vstack(m1)
    m2 = np.vstack(m2)

    write_coil_pos_hdf5(fn_hdf, centers, m0, m1, m2, datanames=datanames, data=data, overwrite=overwrite)


def create_stimsite_from_exp_hdf5(fn_exp, fn_hdf, datanames=None, data=None, overwrite=False):
    """
    This takes an experiment.hdf5 file and creates an .hdf5 + .xdmf tuple for all coil positions for visualization.

    Parameters
    ----------
    fn_exp : str
        Path to experiment.hdf5
    fn_hdf : str
        Filename for the resulting .hdf5 file. The .xdmf is saved with the same basename.
        Folder should already exist.
    datanames : str or list of str, optional
        Dataset names for ``data``
    data : np.ndarray, optional
        Dataset array with shape = ``(len(poslist.pos), len(datanames())``.
    overwrite : bool, default: False
        Overwrite existing files.
    """
    df_stim = pd.read_hdf(fn_exp, "stim_data")

    matsimnibs = np.zeros((4, 4, df_stim.shape[0]))

    for i in range(df_stim.shape[0]):
        matsimnibs[:, :, i] = df_stim["coil_mean"].iloc[i]

    create_stimsite_from_matsimnibs(fn_hdf=fn_hdf,
                                    matsimnibs=matsimnibs,
                                    datanames=datanames,
                                    data=data,
                                    overwrite=overwrite)


def create_stimsite_from_matsimnibs(fn_hdf, matsimnibs, datanames=None, data=None, overwrite=False):
    """
    This takes a matsimnibs array and creates an .hdf5 + .xdmf tuple for all coil positions for visualization.

    Centers and coil orientations are written disk.

    Parameters
    ----------
    fn_hdf: str
        Filename for the .hdf5 file. The .xdmf is saved with the same basename.
        Folder should already exist.
    matsimnibs: np.ndarray
        (4, 4, n_pos)
        Matsimnibs matrices containing the coil orientation (x,y,z) and position (p)

        .. math::
            \\begin{bmatrix}
            | & | & | & | \\\\
            x & y & z & p \\\\
            | & | & | & | \\\\
            0 & 0 & 0 & 1 \\\\
            \\end{bmatrix}
    datanames: str or list of str, optional
        Dataset names for ``data``.
    data: np.ndarray, optional
        (len(poslist.pos), len(datanames).
    overwrite : bool, default: False
        Overwrite existing files.
    """
    matsimnibs = np.atleast_3d(matsimnibs)
    n_pos = matsimnibs.shape[2]
    centers = np.zeros((n_pos, 3))
    m0 = np.zeros((n_pos, 3))
    m1 = np.zeros((n_pos, 3))
    m2 = np.zeros((n_pos, 3))
    if data is not None:
        assert isinstance(data, np.ndarray)

    for i in range(matsimnibs.shape[2]):
        centers[i, :] = matsimnibs[0:3, 3, i]
        m0[i, :] = matsimnibs[0:3, 0, i]
        m1[i, :] = matsimnibs[0:3, 1, i]
        m2[i, :] = matsimnibs[0:3, 2, i]

    write_coil_pos_hdf5(fn_hdf, centers, m0, m1, m2, datanames=datanames, data=data, overwrite=overwrite)


def write_coil_pos_hdf5(fn_hdf, centers, m0, m1, m2, datanames=None, data=None, overwrite=False):
    """
    Creates a ``.hdf5`` + ``.xdmf`` file tuple for all coil positions.
    Coil centers and coil orientations are saved, and - optionally - data for each position if ``data`` and
    ``datanames`` are provided.

    Parameters
    ----------
    fn_hdf : str
        Filename for the .hdf5 file. The .xdmf is saved with the same basename.
        Folder should already exist.
    centers : np.ndarray of float
        (n_pos, 3) Coil positions.
    m0 : np.ndarray of float
        (n_pos, 3) Coil orientation x-axis (looking at the active (patient) side of the coil pointing to the right).
    m1 : np.ndarray of float
        (n_pos, 3) Coil orientation y-axis (looking at the active side of the coil pointing up away from the handle).
    m2 : np.ndarray of float
        (n_pos, 3) Coil orientation z-axis (looking at the active (patient) side of the coil pointing to the patient).
    datanames : str or list of str, optional
        (n_data) Dataset names for ``data``
    data : np.ndarray, optional
        (n_pos, n_data) Dataset array with (len(poslist.pos), len(datanames()).
    overwrite : bool, default: False
        Overwrite existing files.
    """
    n_pos = centers.shape[0]
    if isinstance(datanames, str):
        datanames = [datanames]

    if data is not None:
        if datanames is None:
            raise ValueError("Provide datanames= with data= argument.")
        if isinstance(datanames, str):
            datanames = [datanames]
        if len(data.shape) <= 1:
            data = np.atleast_1d(data)[:, np.newaxis]
        assert data.shape == (n_pos, len(datanames))
    if datanames is not None and data is None:
        raise ValueError("Provide data= with datanames= argument.")

    m0_reshaped = np.hstack((m0, np.zeros((n_pos, 1)))).T[:, np.newaxis, :]
    m1_reshaped = np.hstack((m1, np.zeros((n_pos, 1)))).T[:, np.newaxis, :]
    m2_reshaped = np.hstack((m2, np.zeros((n_pos, 1)))).T[:, np.newaxis, :]
    centers_reshaped = np.hstack((centers, np.ones((n_pos, 1)))).T[:, np.newaxis, :]

    matsimnibs = np.concatenate((m0_reshaped, m1_reshaped, m2_reshaped, centers_reshaped), axis=1)

    # write hdf5 file
    if not fn_hdf.endswith('.hdf5'):
        fn_hdf += '.hdf5'
    if os.path.exists(fn_hdf) and not overwrite:
        raise OSError(fn_hdf + " already exists. Set overwrite flag for create_stimsite_from_poslist.")

    with h5py.File(fn_hdf, 'w') as f:
        f.create_dataset('centers', data=centers.astype(np.float64))
        f.create_dataset('m0', data=m0.astype(np.float64))
        f.create_dataset('m1', data=m1.astype(np.float64))
        f.create_dataset('m2', data=m2.astype(np.float64))
        f.create_dataset("matsimnibs", data=matsimnibs)

        if data is not None:
            for i, col in enumerate(data.T):
                f.create_dataset('/data/' + datanames[i], data=col)

    # write .xdmf file
    with open(fn_hdf[:-4] + 'xdmf', 'w') as f:
        fn_hdf = os.path.basename(fn_hdf)  # relative links

        # header
        f.write('<?xml version="1.0"?>\n')
        f.write('<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>\n')
        f.write('<Xdmf Version="2.0" xmlns:xi="http://www.w3.org/2001/XInclude">\n')
        f.write('<Domain>\n')
        f.write('<Grid CollectionType="Spatial" GridType="Collection" Name="Collection">\n')

        # one grid for coil dipole nodes...store data hdf5.
        #######################################################
        f.write('<Grid Name="stimsites" GridType="Uniform">\n')
        f.write(f'<Topology NumberOfElements="{centers.shape[0]}" TopologyType="Polyvertex" Name="Tri">\n')
        f.write(f'\t<DataItem Format="XML" Dimensions="{centers.shape[0]} 1">\n')
        np.savetxt(f, list(range(centers.shape[0])), fmt='\t%d', delimiter=' ')  # 1 2 3 4 ... N_Points
        f.write('\t</DataItem>\n')
        f.write('</Topology>\n\n')

        # nodes
        f.write('<Geometry GeometryType="XYZ">\n')
        f.write(f'\t<DataItem Format="HDF" Dimensions="{centers.shape[0]} 3">\n')
        f.write(f'\t{fn_hdf}:/centers\n')
        f.write('\t</DataItem>\n')
        f.write('</Geometry>\n\n')

        # data
        # dipole magnitude
        # the 4 vectors
        for i in range(3):
            f.write(f'\t\t<Attribute Name="dir_{i}" AttributeType="Vector" Center="Cell">\n')
            f.write(f'\t\t\t<DataItem Format="HDF" Dimensions="{centers.shape[0]} 3">\n')
            f.write(f'\t\t\t{fn_hdf}:/m{i}\n')
            f.write('\t\t\t</DataItem>\n')
            f.write('\t\t</Attribute>\n\n')

        if data is not None:
            for i, col in enumerate(data.T):
                f.write(f'\t\t<Attribute Name="{datanames[i]}" AttributeType="Scalar" Center="Cell">\n')
                f.write('\t\t\t<DataItem Format="HDF" Dimensions="' + str(centers.shape[0]) + ' 1">\n')
                f.write(f'\t\t\t{fn_hdf}:/data/{datanames[i]}\n')
                f.write('\t\t\t</DataItem>\n')
                f.write('\t\t</Attribute>\n\n')

        f.write('</Grid>\n')
        # end coil dipole data

        # footer
        f.write('</Grid>\n')
        f.write('</Domain>\n')
        f.write('</Xdmf>\n')
