import os
import re
import glob
from collections import defaultdict
import yaml
import logging

logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())


def load_yaml(filename):
    """
    reads data stored in a yaml file to a dict which is returned
    """
    with open(filename, "r") as handle:
        return yaml.safe_load(handle)


def get_yamls_in_path(path=""):
    """
    recursively searches 'path' for files matching '*.yaml'
    :param path: string, path that is searched
    :return: list of strings, found yamls
    """
    return [f for f in glob.glob(os.path.join(path, "**/*.yaml"), recursive=True)]


def write_yaml(dictionary, filename, *args, **kwargs):
    with open(filename, 'w') as handle:
        yaml.dump(dictionary, handle, *args, default_flow_style=False, **kwargs)


def is_existing_file_writeable(f):
    return os.access(f, os.W_OK)


def is_path_writable_or_can_be_created(f):
    if os.path.exists(f):
        return is_existing_file_writeable(f)
    else:
        parent_dir = os.path.dirname(f)
        if parent_dir == f:
            return False
        else:
            return is_path_writable_or_can_be_created(parent_dir)


def is_file_top_dir_writeable(f):
    return os.access(os.path.dirname(f), os.W_OK)


def replace_regex_in_file(fname, find, replace, dry_run=False):
    regex = re.compile(find, flags=re.MULTILINE)
    if not is_existing_file_writeable(fname):
        raise PermissionError(f"Could not manipulate file because it is not writeable:\n\t{fname}")
    with open(fname) as handle:
        content = handle.read()
        content = regex.sub(replace, content)
    if dry_run:
        print("Dry running...")
        print(content)
    else:
        with open(fname, "w") as handle:
            handle.write(content)
    logger.info(fname, "adapted")


def iter_file(fname):
    with open(fname, "r") as handle:
        for line in handle:
            line = line.strip()
            if len(line) > 0 and line[0] != "#":
                yield line.strip("\r\n")


def read_taxmap(taxmap, taxlevel):
    """
    read taxmap into dictionary, only taxlevel information is returned

    :param taxmap: tsv file, mapping accessions to taxa
    :param taxlevel: list of taxlevel names
    :return: dictionary, keys: accessions, values: tax
    """
    taxa = {}
    l = -len(taxlevel)+1
    with open(taxmap, "r") as handle:
        for line in handle:
            line = line.strip()
            if len(line) == 0:
                continue
            fields = line.split("\t")
            taxa[fields[0]] = fields[l:]
    return taxa


def get_headline(fname):
    if not os.path.exists(fname):
        raise FileNotFoundError(f"Could not read headline from file: {fname}")
    with open(fname, "r") as handle:
        return handle.readline().lstrip("#").strip()


def get_accs_from_pg_yaml(pg_yaml):
    """get set containing all accessions stored in the defined protein group yaml file"""
    pgs = load_yaml(pg_yaml)
    sep = pgs['style']['protein_sep']
    accs = set()
    for pg in pgs['protein_groups']:
        accs.update(pg['proteins'].split(sep))
    return accs


def is_with_spectra(pg_yaml):
    protein_groups = load_yaml(pg_yaml)
    with_spectra = True if 'spectra' in protein_groups['protein_groups'][0].keys() else False
    return with_spectra


def get_pg_dict(pg_yaml):
    pgs = load_yaml(pg_yaml)
    return pgs['protein_groups']


def get_samples(pg_yaml):
    pgs = load_yaml(pg_yaml)
    samples = pgs['samples']
    samples.sort()
    return samples


def get_groups_to_align_from_pg_yaml(pg_yaml):
    data = pg_yaml if (type(pg_yaml) is dict) else load_yaml(pg_yaml)
    sep = data['style']['protein_sep']
    pgs = data['protein_groups']
    return [x for x in range(len(pgs)) if sep in pgs[x]['proteins']]


def load_protein_groups_from_yaml(pg_yaml):
    protein_group_yaml = load_yaml(pg_yaml)
    sep = protein_group_yaml['style']['protein_sep']
    protein_groups_dict = protein_group_yaml['protein_groups']
    return protein_groups_dict, sep


def get_accs_per_pg_from_yaml(pg_yaml):
    """
    :params yaml: protein group yaml
    :return pgid_to_accs_dict: key: int protein group number, value = list of associated protein accessions
    """
    protein_groups_dict, sep = load_protein_groups_from_yaml(pg_yaml)
    pgid_to_accs_dict = defaultdict(list)
    for protein_group_number, pg_dict in enumerate(protein_groups_dict):
        pgid_to_accs_dict[protein_group_number] = pg_dict['proteins'].split(sep)
    return pgid_to_accs_dict


def get_spectra_per_acc_from_yaml(pg_yaml):
    """
    :params yaml: protein group yaml
    :return accs_to_spectra_dict: key: string protein accessions value = set of associated spectra IDs
    """
    # protein_groups_dict[protein_group_number]['spectra'] = ['[[spectra1, spectra2]]', '[[spectra1, spectra3]]']
    protein_groups_dict, sep = load_protein_groups_from_yaml(pg_yaml)
    with_spectra = is_with_spectra(pg_yaml)
    accs_to_spectra_dict = defaultdict(set)
    pgid_to_accs_dict = get_accs_per_pg_from_yaml(pg_yaml)
    if with_spectra:
        for protein_group_number in range(len(protein_groups_dict)):
            # spectra_per_pg_list: one list of spectra per pg: '[[spectra1, spectra2]]'
            spectra_per_pg_list = []
            for sample_spectra in protein_groups_dict[protein_group_number]['spectra']:
                # sample_spectra: spectra from one sample in pg '[[spectra1, spectra2][spectra3]]' (one string)
                if sample_spectra == '-':
                    continue
                # change string '[[spectra1, spectra2][spectra1]' to [['spectra1', 'spectra2'],['spectra1']]
                # "[[ms_run[5]:scan=18518][ms_run[10]:scan=13967]]" to [['ms_run5:scan=18518']['ms_run10:scan=13967']]
                spectra_per_pg_per_sample_list = [spectra_str.split(sep) for spectra_str in [spectra_per_protein.replace('[', '')
                                       .replace(']', '') for spectra_per_protein in sample_spectra.split(f"]{sep}[")]]
                # spectra_per_pg_list empty, initialize with first entry
                if not spectra_per_pg_list:
                    spectra_per_pg_list = spectra_per_pg_per_sample_list
                # for all samples, add s
                else:
                    for i, spectra in enumerate(spectra_per_pg_per_sample_list):
                        spectra_per_pg_list[i].extend(spectra)
            # fill accs_to_spectra_dict
            for i, acc in enumerate(pgid_to_accs_dict[protein_group_number]):
                accs_to_spectra_dict[acc] = set(spectra_per_pg_list[i])
    else:
        accs_to_spectra_dict = None
    return accs_to_spectra_dict


def get_spectra_header_and_accs_to_spectraID_dict(pg_yaml):
    """
    param pg_yaml: protein group yaml
    return spectra_header:[string list], spectra_IDs_per_accession_dict[acc]=[spectra]
    """
    protein_group_yaml = load_yaml(pg_yaml)
    protein_groups_dict, sep = load_protein_groups_from_yaml(pg_yaml)
    spectra_header = []
    for sample in protein_group_yaml['samples']:
        spectra_header.extend([sample + ":: spectra_IDs"])
    spectra_IDs_per_accession_dict = {}
    for i in range(len(protein_group_yaml['protein_groups'])):
        for k, acc in enumerate(protein_group_yaml['protein_groups'][i]['proteins'].split(sep)):
            spectra_IDs_per_accession_dict[acc] = []
            for spectra_per_sample in protein_group_yaml['protein_groups'][i]['spectra']:
                spectra_lists = [s.replace('[', '').replace(']', '') for s in spectra_per_sample.split('][')]
                spectra_lists = [';'.join(spectra_list.split(', ')) for spectra_list in spectra_lists]
                try:
                    spectra_IDs_per_accession_dict[acc].append(spectra_lists[k])
                except IndexError:
                    spectra_IDs_per_accession_dict[acc].append('-')

    return spectra_header, spectra_IDs_per_accession_dict