import pyerrors as pe
import datalad.api as dl


def read_param(path, project, file_in_project):
    """
    Read the parameters from the sfcf file.

    Parameters
    ----------
    file : str
        The path to the sfcf file.

    Returns
    -------
    dict
        The parameters in the sfcf file.

    """

    file = path + "/projects/" + project + '/' + file_in_project
    dl.get(file, dataset=path)
    with open(file, 'r') as f:
        lines = f.readlines()

    params = {}
    params['wf_offsets'] = []
    params['wf_basis'] = []
    params['wf_coeff'] = []
    params['qr'] = {}
    params['mrr'] = []
    params['crr'] = []
    params['qs'] = {}
    params['mrs'] = []
    params['crs'] = []

    for line in lines:
        if line.startswith('#'):
            continue
        if line.startswith('\n'):
            continue
        if line.startswith('wf_offsets'):
            num_wf_offsets = line.split()[1]
            for i in range(int(num_wf_offsets)):
                params['wf_offsets'].append([float(x) for x in lines[lines.index(line) + i + 1].split("#")[0].split()])

        if line.startswith('wf_basis'):
            num_wf_basis = line.split()[1]
            for i in range(int(num_wf_basis)):
                params['wf_basis'].append([float(x) for x in lines[lines.index(line) + i + 1].split("#")[0].split()])

        if line.startswith('wf_coeff'):
            num_wf_coeff = line.split()[1]
            for i in range(int(num_wf_coeff)):
                params['wf_coeff'].append([float(x) for x in lines[lines.index(line) + i + 1].split("#")[0].split()])

        if line.startswith('qr'):
            num_qr = line.split()[1]
            for i in range(int(num_qr)):
                dat = lines[lines.index(line) + i + 1].split("#")[0].strip().split()[:-1]
                params['qr'][dat[0]] = {}
                params['qr'][dat[0]]['mass'] = float(dat[1])
                params['qr'][dat[0]]['thetas'] = [float(x) for x in dat[2:5]]

        if line.startswith('mrr'):
            num_mrr = line.split()[1]
            for i in range(int(num_mrr)):
                params['mrr'].append(lines[lines.index(line) + i + 1].split("#")[0].strip())

        if line.startswith('crr'):
            num_crr = line.split()[1]
            for i in range(int(num_crr)):
                params['crr'].append(lines[lines.index(line) + i + 1].split("#")[0].strip())

        if line.startswith('qs'):
            num_qs = line.split()[1]
            for i in range(int(num_qs)):
                dat = lines[lines.index(line) + i + 1].split("#")[0].strip().split()[:-1]
                params['qs'][dat[0]] = {}
                params['qs'][dat[0]]['mass'] = float(dat[1])
                params['qs'][dat[0]]['thetas'] = [float(x) for x in dat[2:5]]

        if line.startswith('mrs'):
            num_mrs = line.split()[1]
            for i in range(int(num_mrs)):
                params['mrs'].append(lines[lines.index(line) + i + 1].split("#")[0].strip())

        if line.startswith('crs'):
            num_crs = line.split()[1]
            for i in range(int(num_crs)):
                params['mrs'].append(lines[lines.index(line) + i + 1].split("#")[0].strip())

    # catch standard cases
    if params['wf_offsets'] == []:
        params['wf_offsets'] = [[0, 0, 0]]

    if params['wf_coeff'] == []:
        params['wf_coeff'] = [[0, -1]]

    return params


def _map_params(params, spec_list):
    """
    Map the extracted parameters to the extracted data.

    """

    new_specs = {}
    # quarks
    quarks = spec_list[0].split(" ")

    new_specs['quarks'] = (params['qr'][quarks[0]], params['qr'][quarks[1]])

    # offset
    new_specs['offset'] = (params['wf_offsets'][int(spec_list[1])])
    # wf1
    contribs = []
    for i, coeff in enumerate(params['wf_coeff'][int(spec_list[2])]):
        if not coeff == 0:
            contrib = (coeff, params['wf_basis'][i])
            contribs.append(contrib)
    new_specs['wf1'] = contribs

    if len(spec_list) == 4:
        # wf2
        contribs = []
        for i, coeff in enumerate(params['wf_coeff'][int(spec_list[3])]):
            if not coeff == 0:
                contrib = (coeff, params['wf_basis'][i])
                contribs.append(contrib)
        new_specs['wf2'] = contribs
    return new_specs


def read_data(path, project, dir_in_project, prefix, param, version='1.0c', cfg_seperator='n', sep='/', **kwargs):
    """
    Extract the data from the sfcf file.

    Returns
    -------
    dict
        The data from the sfcf file.

    """
    names = kwargs.get('names', None)
    corr_types = {
        'f_A': 'bi',
        'f_P': 'bi',
        'g_A': 'bi',
        'g_P': 'bi',
        'f_1': 'bb',
        'k_1': 'bb',
    }
    directory = path + "/projects/" + project + '/' + dir_in_project
    dl.get(directory, dataset=path)
    corr_type_list = []
    for corr_name in param['crr']:
        if corr_name not in corr_types:
            raise ValueError('Correlator type not known.')
        corr_type_list.append(corr_types[corr_name])

    if not param['crr'] == []:
        if names is not None:
            data_crr = pe.input.sfcf.read_sfcf_multi(directory, prefix, param['crr'], param['mrr'], corr_type_list, range(len(param['wf_offsets'])),
                                                     range(len(param['wf_basis'])), range(len(param['wf_basis'])), version, cfg_seperator, keyed_out=True, names=names)
        else:
            data_crr = pe.input.sfcf.read_sfcf_multi(directory, prefix, param['crr'], param['mrr'], corr_type_list, range(len(param['wf_offsets'])),
                                                     range(len(param['wf_basis'])), range(len(param['wf_basis'])), version, cfg_seperator, keyed_out=True)

    if not param['crs'] == []:
        data_crs = pe.input.sfcf.read_sfcf_multi(directory, param['crs'])

    data = {}
    if not param['crr'] == []:
        for key in data_crr.keys():
            data[key] = data_crr[key]

    if not param['crs'] == []:
        for key in data_crs.keys():
            data[key] = data_crs[key]

    # sort data by correlator
    sorted_data = {}
    for key in data.keys():
        key_parts = key.split(sep)
        corr = key_parts[0]
        if corr_types[corr] == 'bi':
            specs = _map_params(param, key_parts[1:-1])
        else:
            specs = _map_params(param, key_parts[1:])
        if corr not in sorted_data:
            sorted_data[corr] = {}
        sorted_data[corr][sep.join(key_parts[1:])] = data[key]
    return sorted_data, specs