#! /usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import os
import numpy as np
import pandas as pd
from pathlib import Path
from GRATIOSA.useful_functions_Chipseq import *
from GRATIOSA.globvar import *
from GRATIOSA.Genome import Genome

class Chipseq:
    '''
    The Chipseq class is mainly used to collect coverage data (Chipseq signals) 
    along a genome. The associated methods and functions allow the binning and 
    averaging of these signals. The enrichment peaks positions can also be loaded 
    as an attribute of this class. 
    Note: this class is only implemented for single-chromosome genomes with an annotation. 

    Each Chipseq instance has to be initialized with an organism name
    Example:
        >>> from GRATIOSA import Chipseq
        >>> gen = Genome.Genome("dickeya")
        >>> ch = Chipseq.Chipseq(gen)
    '''

    def __init__(self, gen):

        if not hasattr(gen, "genes"):
            print("Initializing a Chipseq object requires an annotated genome. Trying to load annotation.")
            gen.load_annotation()
        
        self.genome = gen
        self.name = gen.name
        self.genes = gen.genes
        # for multiple chromosomes,
        # create one chipseq object for each chromosomes
        # accessed through self.chipseqs
        if not hasattr(gen,"contig"):
            print("ERROR: a sequence and annotation must be loaded before loading chipseq")
        if gen.contig:
            self.chipseqs = self
        else:
            self.chipseqs = []
            for ge in gen.genomes:
                ch = Chipseq(ge)
                self.chipseqs.append(ch)


    def load_signal(self, cond="all"):
        '''
        load_signal imports a 1D distribution along the chromosome (typically a 
        CHIPSeq distribution) from a data file (typically a .bedGraph obtained 
        with bamCompare) containing: (1) bin starting positions, (2) bin ending 
        positions and (3) signal in  each bin.

        Creates a new attribute of the Chipseq instance: 
            * self.signal (dict.)
                   Dictionary of shape {condition: array containing one signal
                   value per genomic position}

        Args:
            cond (list of str.): 
                    Selection of one or several conditions (1 condition 
                    corresponds to 1 data file). By default: cond ='all' ie all 
                    available signals are loaded. All selected conditions have to 
                    be listed in signal.info file.

        Note:
            The data importation requires an signal.info file that contains the 
            columns positions of each information in the data file, in the following 
            order:
            [0] condition, [1] filename, [2] separator used in the data file, 
            [3] bin_start, [4] bin_end, [5] signal, 
            [6] chromosome name (def 0, not used for single-chromosome genome)

        Example:
            >>> from GRATIOSA import Chipseq
            >>> ch = Chipseq.Chipseq("ecoli")
            >>> ch.load_signal()
            >>> ch.signal['Signal_Test']
            array([100,   0,   0, ..., 100, 100,   0])
        '''
        # gets the path to data and .info files
        path2dir = f"{basedir}data/{self.name}/chipseq/signals/"

        # converts from str to list (useful if only 1 condition is selected)
        if isinstance(cond, str):
            cond = [cond]

        # tries to open .info file
        list_cond_info = []  
        if Path(f"{path2dir}signals.info").exists():
            with open(f"{path2dir}signals.info", "r") as f:
                header = next(f)

                # 1 line in the .info file corresponds to 1 condition
                # (1 data file)
                for line in f:
                    line = line.strip().split('\t')
                    list_cond_info.append(line[0])

                    # For each selected condition:
                    if cond == ["all"] or line[0] in cond:
                        print('Loading condition', line[0])
                        # loads data
                        separator = line[2]
                        if separator == '\\t':
                            separator = '\t'
                        data = pd.read_csv(
                            path2dir + line[1], sep=separator, header=None)
  
                        if len(line) == 7:
                            chrom = int(line[6])
                        else:
                            chrom = 0
                        chromosomes = data.iloc[:, chrom][0]

                        # Separate chromosomes
                        if self.genome.contig:
                            datas=[data]
                            objs=[self]
                        else:
                            datas=[data[chromosomes==x] for x in self.genome.chromosome_name]
                            objs=self.chipseqs

                        # attribute values to the relevant chipseq object
                        for id,d in enumerate(datas):
                            # computes binsize of each bin in the file
                            signal = d.iloc[:, int(line[5])]
                            bs = d.iloc[:, int(line[4])] - \
                                d.iloc[:, int(line[3])]
                            # creates a new attribute to the Chipseq instance:
                            # signal[cond]= signal per pos
                            ch = objs[id]
                            if not hasattr(ch, "signal"):
                                ch.signal = {}
                            ch.signal[line[0]] = np.array(np.repeat(signal, bs))
            f.close()

        else:
            print(f"Unable to locate signals.info in {path2dir}")

        # warns that some selected conditions aren't described in the
        # .info file
        for c in cond:
            if c not in list_cond_info and c != "all":
                print(f"Please add {c} in signals.info")


    def load_binned_signal(self, binsize, cond="all"):
        '''
        If a file containing the data for the chosen condition and binsize 
        exists, load_binned_signal loads these data using the load_signal method
        Else, load_binned signal performs the following steps: 
        1 - imports a 1D distribution along the chromosome (typically a CHIPSeq 
        distribution) using the load_signal method
        2 - performs the binning at the chosen binsize using the 
        useful_functions_Chipseq.binning
        3 - saves the binned data in a file 

        Creates or adds items to 2 Chipseq instance attributes:
            * self.signal (dict.)
                    Dictionary of shape {condition: array containing one signal
                    value per genomic position (before binning)}
            * self.binned_signal (dict.)
                    Dictionary of shape {cond_bin: array containing one binned 
                    signal value per genomic position} with cond_bin the condition 
                    name merged with the bin size (example: WT_bin200b). 

        Args:
            binsize (int.)
            cond (list of str.): 
                    Selection of one or several conditions (1 condition 
                    corresponds to 1 data file). By default: cond ='all' ie all 
                    available signals are loaded. All selected conditions have 
                    to be listed in signal.info file.

        Note:
            See ChipSeq.load_signal method for the data requirements.

        Example:
            >>> from GRATIOSA import Chipseq
            >>> ch = Chipseq.Chipseq("ecoli")
            >>> ch.load_binned_signal(binsize=100,cond='Signal_Test')
            >>> ch.binned_signal["Signal_Test_bin100b"]
            array([100,   100,   100, ..., 10, 10, 10])
        '''
        if not self.genome.contig:
            print("Error: only implemented for single-chromosome genome objects.")
            return 1
        
        if not hasattr(self, "length"):
            gen = Genome(self.name)
            gen.load_seq()
            self.length = gen.length

        # gets the path to data and .info files
        path2files = f"{basedir}data/{self.name}/chipseq/signals/"
        f_path = f"{path2files}/binned_data/"

        # gets the list of all existing conditions using .info file
        if Path(path2files + 'signals.info').exists():
            if cond == "all":
                with open(path2files + 'signals.info', "r") as fi:
                    header = next(fi)
                    cond = []
                    for line in fi:  # for each condition
                        line = line.strip().split('\t')
                        cond.append(line[0])
                fi.close()

            # converts cond from str to list if only 1 condition is selected
            elif isinstance(cond, str):
                cond = [cond]

        else:
            print(f"Unable to locate signals.info in {path2files}")

        # For each selected condition:
        if not hasattr(self, "binned_signal"):
            self.binned_signal = {}  # creates the new attribute

        for c in cond:
            cond_name = f"{c}_bin{str(binsize)}b"
            print('Condition:' + cond_name)

            # if a file contains data for this condition and binsize
            if Path(f"{f_path}{cond_name}.npy").exists():
                # adds the data to the binned_signal attribute
                print(f"loading existing file: {cond_name}.npy")
                binned_data = np.load(f"{f_path}{cond_name}.npy")
                self.binned_signal[cond_name] = np.array(
                    np.repeat(binned_data, binsize))[:self.length]

            # if no file contains data for this condition and binsize
            else:
                print("performing the binning")
                # loads data
                self.load_signal(cond=c)
                # performs the binning using the
                # useful_functions_Chipseq.binning function
                binned_data = binning(self.signal[c], binsize=binsize,
                                      stat="mean").statistic
                # adds the data to the binned_signal attribute
                self.binned_signal[cond_name] = np.array(
                    np.repeat(binned_data, binsize))[:self.length]

                # creates the "binned_data" directory if necessary
                Path(f_path).mkdir(parents=True, exist_ok=True)
                # saves the data as a .npy file
                np.save(f"{f_path}{cond_name}.npy", binned_data)

    def load_smoothed_signal(self, window, cond="all", *args, **kwargs):
        '''
        If a data file containing the data for the chosen condition and 
        smoothing exists, load_binned_signal loads these data using the 
        load_signal method
        Else, load_smoothed_signal performs the following steps: 
        1 - imports a 1D distribution along the chromosome (typically 
        a CHIPSeq distribution) using the load_signal method
        2 - performs the smoothing (moving average with the chosen 
        window size) using the useful_functions_Chipseq.smoothing
        3 - saves the smoothed data in a file 
        
        Creates or adds items to 2 Chipseq instance attributes:
            * self.signal (dict.)
                   Dictionary of shape {condition: array containing one signal
                   value per genomic position (before smoothing)}
            * self.smoothed_signal (dict.)
                    Dictionary of shape {cond_smoo: array containing one binned 
                    signal value per genomic position} with cond_smoo the 
                    condition name merged with the smoothing window size 
                    (example:  WT_smooth200b).


        Args:
            window (int.): window size. The value of the smoothed signal 
                    of a position p is equal to the average of the signal 
                    between p - window/2 and p + window/2.
            cond (list of str.): 
                    Selection of one or several conditions (1 condition 
                    corresponds to 1 data file). By default: cond ='all' ie 
                    all available signals are loaded. All selected conditions 
                    have to be listed in signal.info file.
        
        Note:
            See Chipseq.load_signal method for the data requirements.

        Example:
            >>> from GRATIOSA import Chipseq
            >>> ch = Chipseq.Chipseq("ecoli")
            >>> ch.load_smoothed_signal(window=100,cond='Signal_Test')
            >>> ch.smoothed_signal["Signal_Test_smooth100b"]
            array([100.1,   99.8,   98.8, ..., 10.1, 11.1, 10.8])
        '''

        if not self.genome.contig:
            print("Error: only implemented for single-chromosome genome objects.")
            return 1

        # gets the path to data and .info files
        path2files = f"{basedir}data/{self.name}/chipseq/signals/"
        f_path = f"{path2files}/smoothed_data/"

        # gets the list of all existing conditions using .info file
        info_filename = kwargs.get('info_filename', 'signals.info')
        if Path(path2files + info_filename).exists():
            if cond == "all":
                with open(path2files + info_filename, "r") as fi:
                    header = next(fi)
                    cond = []
                    for line in fi:  # for each condition
                        line = line.strip().split('\t')
                        cond.append(line[0])
                fi.close()
            # converts cond from str to list if only 1 condition is selected
            elif isinstance(cond, str):
                cond = [cond]
        else:
            print("Unable to locate " + info_filename)

        if not hasattr(self, "smoothed_signal"):
            self.smoothed_signal = {}  # creates the new attribute

        # For each selected condition:
        for c in cond:
            cond_name = f"{c}_smooth{str(window)}b"
            print('Condition:' + cond_name)

            # if a file contains data for this condition and window size
            if Path(f"{f_path}{cond_name}.npy").exists():
                # adds the data to the smoothed_signal attribute
                print(f"loading existing file: {cond_name}.npy")
                self.smoothed_signal[cond_name] = np.load(
                    f"{f_path}{cond_name}.npy")

            # if no file contains data for this condition and binsize
            else:
                print("performing the smoothing")
                # loads data
                self.load_signal(cond=c)
                # performs the smoothing using the
                # useful_functions_Chipseq.smoothing function
                smooth_data = smoothing(
                    self.signal[c], window=window)
                # adds the data to the smoothed_signal attribute
                self.smoothed_signal[cond_name] = smooth_data
                # creates the "smoothed_data" directory if necessary
                Path(f_path).mkdir(parents=True, exist_ok=True)
                # saves the data as a .npy file
                np.save(f"{f_path}{cond_name}.npy", smooth_data)


    def load_signals_average(self, list_cond, average_name, *args, **kwargs):
        '''
        Load_signals_average computes and loads the average of signals replicates.
        First, the function loads the signals (which must be listed in the 
        signals.info file in the /chipseq/signals/ folder) using the load_signal 
        method. It can then process the data using the load_binned_signal or 
        load_smoothed_signal methods. Finally, the average of these signals at 
        each genomic position of the genome is calculated. This average signal 
        is assigned to the Chipseq instance as signals_average attribute:
        
        * self.signals_average (dict.):
                Dictionary of shape {average_name: array containing one 
                averaged signal value per genomic position}. 

        Args:
            list_cond (list of str.): 
                    Selection of conditions that will be averaged. All selected 
                    conditions have to be listed in signal.info file.
            average_name (str.): 
                    Name of the obtained signal
            data_treatment (Optionnal [str.] "binning", "smoothing" or None):
                    Treatment to be applied to the different signals 
                    before averaging them (None by default).
            window (Optionnal [int.]): 
                    window size used only if data_treatment = "smoothing"
                    The value of the smoothed signal of a position p is 
                    equal to the average of the signal between 
                    p - window/2 and p + window/2
            binsize (Optionnal [int.]): bin size used only if data_treatment = "binning"
        
        Note:
            See Chipseq.load_signal method for the data requirements

        Example:
            >>> from GRATIOSA import Chipseq
            >>> ch = Chipseq.Chipseq("ecoli")
            >>> ch.load_signals_average(list_cond=["Signal1","Signal2"],
            ...                         average_name="Mean_signal",
            ...                         data_treatment = "smoothing",
            ...                         window=500)
            >>> ch.signals_average["Mean_signal"]
            array([100.1,   99.8,   98.8, ..., 10.1, 11.1, 10.8])
        '''

        if not self.genome.contig:
            print("Error: only implemented for single-chromosome genome objects.")
            return 1

        
        f_path = f"{basedir}data/{self.name}/chipseq/signals/average_data/"
        data_treatment = kwargs.get('data_treatment', None)
        if not hasattr(self, "signals_average"):
            self.signals_average = {}
        # test if a file for this condition (cond_name) and this binsize
        # already exists

        if Path(f"{f_path}{average_name}.npy").exists():
            with open(f"{f_path}data.info", "r") as f:
                header = next(f).strip("\n").split("\t")
                # 1 line in the .info file corresponds to 1 condition
                for line in f:
                    line = line.strip().split('\t')
                    if line[0] == average_name:
                        print("loading the file obtained with the following parameters:\n")
                        for i in np.arange(len(header)):
                            print(f"{header[i]}: {line[i]}\n")

                        print("Please change 'average_name' to use other signals")
            self.signals_average[average_name] = np.load(
                f"{f_path}{average_name}.npy")

        else:
            print("performing the average")

            if data_treatment == "binning":
                print("data_treatment: binning")
                try:
                    size = kwargs.get('binsize')
                except:
                    sys.exit("please give a binsize")
                self.load_binned_signal(size, cond=list_cond)
                data_to_av = [
                    self.binned_signal[f"{i}_bin{str(size)}b"] for i in list_cond]

            elif data_treatment == "smoothing":
                print("data_treatment: smoothing")
                try:
                    size = kwargs.get('window')
                except:
                    sys.exit(
                        "please give a window size using \"window\" argument")
                self.load_smoothed_signal(size, cond=list_cond)
                data_to_av = [
                    self.smoothed_signal[f"{i}_smooth{str(size)}b"] for i in list_cond]
            else:
                data_treatment = "No treatment"
                size = "NA"
                print("data_treatment: No treatment")
                self.load_signal(cond=list_cond)
                data_to_av = [self.signal[i] for i in list_cond]

            # performs the average between replicates
            average = np.mean(data_to_av, axis=0)

            # adds the signals_average attribute to Chipseq instance
            self.signals_average[average_name] = average

            # saves the data as a .npy file
            Path(f_path).mkdir(parents=True, exist_ok=True)
            np.save(f"{f_path}{average_name}.npy", average)

            # saves information about the used replicates and treatment
            if not Path(f"{f_path}data.info").exists():
                file = open(f"{f_path}data.info", 'w')
                file.write('Name\t Replicates\t Data treatment\t Size (window or bin, in b)')
                file.close()
            file = open(f"{f_path}data.info", 'a')
            file.write(f"\n{average_name}\t{list_cond}\t{data_treatment}\t{size}")
            file.close()


    def get_all_signals(self):
        '''
        get_all_signals groups all loaded signals into a single attribute:

            * self.all_signals (dict.): 
                   Dictionary of shape {condition: array containing one signal
                   value per genomic position}
        '''
        if not self.genome.contig:
            print("Error: only implemented for single-chromosome genome objects.")
            return 1

        d = {}
        for attr in ("signal","binned_signal","smoothed_signal","signals_average"):
            if hasattr(self, attr):
                d.update(getattr(self, attr))
        self.all_signals = d


    def load_signal_per_genes(self, cond='all', window=0):
        """ 
        Load_signal_per_genes computes the mean signal for each Gene (i.e. 
        the mean signal between the gene start and gene end). Genomic signals
        have to be loaded (using load_signal,load_binned_signal, 
        load_smoothed_signal or load_signals_average methods) before using 
        load_signal_per_genes.

        Creates: 
            * self.genes[locus].signal (float.): 
                    new attribute of Gene instances 
                    related to the Chipseq instance given as argument.Contains 
                    the Gene mean signal. 
            * self.signals_gene (dict. of dict.): 
                    Dictionary containing one  subdictionary per condition given 
                    in input. Each  subdictionary contains the signal (value) for 
                    each gene (locus_tag as key). 

        Args:
            cond (Optional [list of str.]): selection of one or several 
                    conditions (1 condition corresponds to 1 data file).
                    By default cond ='all' ie all available loaded signals are 
                    used.
            window (Optional int.): to include the signal around the gene. 
                    The mean signal will be computed between
                    gene start - window and gene end + window.
                    By default, window = 0

        Warning:
            This method needs a genomic annotation. If no annotation is 
            loaded, the load_annotation method with the default "sequence.gff3" 
            file is computed. To use another annotation, please load an annotation 
            to your Transcriptome instance with the following commands before using 
            this method:
            >>> from GRATIOSA import Genome, Chipseq
            >>> g = Genome.Genome(ch.name)
            >>> g.load_annotation(annot_file=chosen_file)
            >>> ch = Chipseq.Chipseq(g)

        Example:
            >>> from GRATIOSA import Chipseq
            >>> ch = Chipseq.Chipseq(g)
            >>> ch.load_signal()
            >>> ch.load_signal_per_genes()
            >>> ch.signals_gene["WT"]["b0984"]
            0.40348639903919425
            >>> ch.genes["b0984"].signal
            {'WT': 0.40348639903919425,
             'Signal_Test': 0.210505235030067}
        """
        self.get_all_signals()
        #if not hasattr(self, "signals_gene"):
        #    self.signals_gene = {}
        #if not hasattr(self, "genes"):
        #    gen = Genome(self.name)
        #    gen.load_annotation()
        #    self.genes = gen.genes
        if cond == "all":
            cond = list(self.all_signals.keys())
        elif isinstance(cond, str):
            cond = [cond]
        
        if not hasattr(self, "signals_gene"):
            self.signals_gene = {}
            for cond_name in cond:
                self.signals_gene[cond_name]={}
            
        for c in cond:
            self.signals_gene[c] = {}
            for locus in self.genes.keys():
                g = self.genes[locus]
                s = np.mean(self.all_signals[c][g.left-window:g.right+window])
                self.genes[locus].add_signal(c, s)
                self.signals_gene[c][locus] = s
        print(f"Gene signals successfully computed for: {cond}")


    def load_peaks(self):
        """ 
        load_peaks imports a list of peaks from a data file (typically 
        a .BED file of peaks obtained with MACS2).

        Creates:
            * self.peaks (dict. of dict.) : new attribut of the Chipseq instance. 
                    dictionary containing one subdictionary per condition 
                    with the shape {(start,end):value}. The key of each
                    subdictionary is a tupe containing the start and end 
                    positions of one peak. The value can for example be the 
                    peak score or the peak height.

        Note:
            The data importation requires a peaks.info file that contains the
            column indices of each information in the data file and some 
            additional information, in the following order:
            [0] Condition, [1] Filename, [2] Startline, 
            [3] Separator, [4] StartCol, [5] StopCol
            [6] Peak value, [7] Chromosome name (def 0)
            peaks.info and data file have to be in the /chipseq/peaks/ directory 

        Example:
            >>> from GRATIOSA import Chipseq
            >>> ch = Chipseq.Chipseq("ecoli")
            >>> ch.load_peaks()
            >>> ch.peaks['test']
            {(840081, 840400): 4.89872,
            (919419, 919785): 5.85158,
            (937220, 937483): 4.87632,
            ...}
        """
        gen = self.genome
        
        # gets the path to data and .info files
        path2dir = f"{basedir}data/{self.name}/chipseq/peaks/"

        # tries to open peaks.info file
        if Path(f"{path2dir}peaks.info").exists():
            with open(f"{path2dir}peaks.info", "r") as f:
                skiphead = next(f)
                # 1 line in peaks.info corresponds to 1 condition (1 data file)
                for line in f:
                    line = line.strip('\n').split('\t')

                    # loads data file information for this condition
                    cond = line[0]
                    path2file = f"{path2dir}{line[1]}"
                    startline, sep = int(line[2]), line[3]
                    start_col, end_col = int(line[4]), int(line[5])
                    val_col = int(line[6])
                    if len(line) == 8:
                        chrom = int(line[6])
                    else:
                        chrom = 0
                    # creates the new attribute using the load_sites_cond function
                    # from useful_functions_Chipseq
                    peaks = load_sites_cond(gen.chromosome_name, path2file, 
                                                       startline, # starting with 1!
                                                       sep,
                                                       start_col, 
                                                       end_col,
                                                       val_col,
                                                       chrom)
                    if gen.contig:
                        # one chromosome
                        # dont even check if the name is right!
                        if not hasattr(self, "peaks"):
                            self.peaks = {}  # creates the new attribute
                        self.peaks[cond] = peaks[gen.chromosome_name]
                    else:
                        for ich,ch in enumerate(self.chipseqs):
                            if not hasattr(ch, "peaks"):
                                ch.peaks={}
                            ch.peaks[cond] = peaks[gen.chromosome_name[ich]]
                    print("Condition %s loaded successfully"%cond)
        else:
            print("No peaks.info, unable to load peaks")
