import pyerrors.input.openQCD as input
import datalad.api as dl
import os
import fnmatch
from typing import Any, Optional


def read_ms1_param(path: str, project: str, file_in_project: str) -> dict[str, Any]:
    """
    Read the parameters for ms1 measurements from a parameter file in the project.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    project: str
        The project from which to read the parameter file.
    file_in_project: str
        The path to the parameter file within the project.

    Returns
    -------
    param: dict[str, Any]
        The parameters read from the file.
    """

    file = os.path.join(path, "projects", project, file_in_project)
    ds = os.path.join(path, "projects", project)
    dl.get(file, dataset=ds)
    with open(file, 'r') as fp:
        lines = fp.readlines()
        fp.close()
    param: dict[str, Any] = {}
    param['rw_fcts'] = []
    param['rand'] = {}

    for i, line in enumerate(lines):
        if line.startswith('#'):
            continue
        if line.startswith('\n'):
            continue
        if line.startswith("nrw"):
            stline = line.strip()
            num_rw = int(stline.split('#')[0].split()[1].strip())
    param["rw_fcts"] = [{}]*num_rw
    for i, line in enumerate(lines):
        for nrw in range(num_rw):
            param["rw_fcts"][nrw] = {}
            if line == "[Reweighting factor " + str(nrw) + "]\n":
                j=0
                while not (lines[i+j] == "\n" or lines[i+j].startswith("[")):
                    j+=1
                chunk = lines[i:j+i]
                for cline in chunk:
                    cline = cline.strip()
                    for rwp in ["rwfact", "nsrc", "irp", "mu", "np"]:
                        if cline.startswith(rwp):
                            param["rw_fcts"][nrw][rwp] = cline.split()[1]
    for nrw in range(num_rw):
        if "nsrc" not in param["rw_fcts"][nrw]:
            param["rw_fcts"][nrw]["nsrc"] = 1
        if "mu" not in param["rw_fcts"][nrw]:
            param["rw_fcts"][nrw]["mu"] = "None"
        if "np" not in param["rw_fcts"][nrw]:
            param["rw_fcts"][nrw]["np"] = "None"
        if "irp" not in param["rw_fcts"][nrw]:
            param["rw_fcts"][nrw]["irp"] = "None"

    return param


def read_ms3_param(path: str, project: str, file_in_project: str) -> dict[str, Any]:
    """
    Read the parameters for ms3 measurements from a parameter file in the project.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    project: str
        The project from which to read the parameter file.
    file_in_project: str
        The path to the parameter file within the project.

    Returns
    -------
    param: dict[str, Any]
        The parameters read from the file.
    """

    file = os.path.join(path, "projects", project, file_in_project)
    ds = os.path.join(path, "projects", project)
    dl.get(file, dataset=ds)
    with open(file, 'r') as fp:
        lines = fp.readlines()
        fp.close()
    param = {}
    for line in lines:
        line = line.strip()
        for rwp in ["integrator", "eps", "ntot", "dnms"]:
            if line.startswith(rwp):
                param[rwp] = line.split()[1]
    return param


def read_rwms(path: str, project: str, dir_in_project: str, param: dict[str, Any], prefix: str, postfix: str="ms1", version: str='2.0', names: Optional[list[str]]=None, files: Optional[list[str]]=None) -> dict[str, Any]:
    """
    Read reweighting factor measurements from the project.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    project: str
        The project from which to read the measurements.
    dir_in_project: str
        The directory within the project where the measurements are located.
    param: dict[str, Any]
        The parameters for the measurements.
    prefix: str
        The prefix of the measurement files.
    postfix: str
        The postfix of the measurement files.
    version: str
        The version of the openQCD used.
    names: list[str]
        Specific names for the replica of the ensemble the measurement file belongs to.
    files: list[str]
        Specific files to read.

    Returns
    -------
    rw_dict: dict[str, dict[str, Any]]
        The reweighting factor measurements read from the files.
    """

    dataset = os.path.join(path, "projects", project)
    directory = os.path.join(dataset, dir_in_project)
    if files is None:
        files = []
        for root, ds, fs in os.walk(directory):
            for f in fs:
                if fnmatch.fnmatch(f, prefix + "*" + postfix + ".dat"):
                    files.append(f)
    dl.get([os.path.join(directory, f) for f in files], dataset=dataset)
    kwargs = {}
    if names is not None:
        kwargs['names'] = names
    if files is not None:
        kwargs['files'] = files
    rwms = input.read_rwms(directory, prefix, version, **kwargs)
    rw_dict: dict[str, dict[str, Any]] = {}
    rw_dict[param["type"]] = {}
    for i in range(len(param["rw_fcts"])):
        par_list= []
        for k in param["rw_fcts"][i].keys():
            par_list.append(str(param["rw_fcts"][i][k]))
        pars = "/".join(par_list)
        rw_dict[param["type"]][pars] = rwms[i]
    return rw_dict


def extract_t0(path: str, project: str, dir_in_project: str, param: dict[str, Any], prefix: str, dtr_read: int, xmin: int, spatial_extent: int, fit_range: int = 5, postfix: str="", names: Optional[list[str]]=None, files: Optional[list[str]]=None) -> dict[str, Any]:
    """
    Extract t0 measurements from the project.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    project: str
        The project from which to read the measurements.
    dir_in_project: str
        The directory within the project where the measurements are located.
    param: dict[str, Any]
        The parameters for the measurements.
    prefix: str
        The prefix of the measurement files.
    dtr_read: int
        The dtr_read parameter for the extraction.
    xmin: int
        The xmin parameter for the extraction.
    spatial_extent: int
        The spatial_extent parameter for the extraction.
    fit_range: int
        The fit_range parameter for the extraction.
    postfix: str
        The postfix of the measurement files.
    names: list[str]
        Specific names for the replica of the ensemble the measurement file belongs to.
    files: list[str]
        Specific files to read.
    """

    dataset = os.path.join(path, "projects", project)
    directory = os.path.join(dataset, dir_in_project)
    if files is None:
        files = []
        for root, ds, fs in os.walk(directory):
            for f in fs:
                if fnmatch.fnmatch(f, prefix + "*" + postfix + ".dat"):
                    files.append(f)
    dl.get([os.path.join(directory, f) for f in files], dataset=dataset)
    kwargs: dict[str, Any] = {}
    if names is not None:
        kwargs['names'] = names
    if files is not None:
        kwargs['files'] = files
    if postfix is not None:
        kwargs['postfix'] = postfix
    kwargs['plot_fit'] = False

    t0 = input.extract_t0(directory,
                            prefix,
                            dtr_read,
                            xmin,
                            spatial_extent,
                            fit_range=fit_range,
                            c=0.3,
                            **kwargs
                            )
    par_list= []
    for k in ["integrator", "eps", "ntot", "dnms"]:
        par_list.append(str(param[k]))
    pars = "/".join(par_list)
    t0_dict: dict[str, Any] = {}
    t0_dict[param["type"]] = {}
    t0_dict[param["type"]][pars] = t0
    return t0_dict


def extract_t1(path: str, project: str, dir_in_project: str, param: dict[str, Any], prefix: str, dtr_read: int, xmin: int, spatial_extent: int, fit_range: int = 5, postfix: str = "", names: Optional[list[str]]=None, files: Optional[list[str]]=None) -> dict[str, Any]:
    """
    Extract t1 measurements from the project.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    project: str
        The project from which to read the measurements.
    dir_in_project: str
        The directory within the project where the measurements are located.
    param: dict[str, Any]
        The parameters for the measurements.
    prefix: str
        The prefix of the measurement files.
    dtr_read: int
        The dtr_read parameter for the extraction.
    xmin: int
        The xmin parameter for the extraction.
    spatial_extent: int
        The spatial_extent parameter for the extraction.
    fit_range: int
        The fit_range parameter for the extraction.
    postfix: str
        The postfix of the measurement files.
    names: list[str]
        Specific names for the replica of the ensemble the measurement file belongs to.
    files: list[str]
        Specific files to read.
    """

    directory = os.path.join(path, "projects", project, dir_in_project)
    if files is None:
        files = []
        for root, ds, fs in os.walk(directory):
            for f in fs:
                if fnmatch.fnmatch(f, prefix + "*" + postfix + ".dat"):
                    files.append(f)
    kwargs: dict[str, Any] = {}
    if names is not None:
        kwargs['names'] = names
    if files is not None:
        kwargs['files'] = files
    if postfix is not None:
        kwargs['postfix'] = postfix
    kwargs['plot_fit'] = False
    t0 = input.extract_t0(directory,
                            prefix,
                            dtr_read,
                            xmin,
                            spatial_extent,
                            fit_range=fit_range,
                            c=2./3,
                            **kwargs
                            )
    par_list= []
    for k in ["integrator", "eps", "ntot", "dnms"]:
        par_list.append(str(param[k]))
    pars = "/".join(par_list)
    t0_dict: dict[str, Any] = {}
    t0_dict[param["type"]] = {}
    t0_dict[param["type"]][pars] = t0
    return t0_dict