#!/usr/bin/env python
# coding: utf-8
from functools import lru_cache
from collections import Counter, defaultdict
import re
import datetime
import torch
import numpy as np
import pickle
from pathlib import Path
from enum import IntEnum
from matplotlib import pyplot as plt
from itertools import product
import pandas as pd
import matplotlib
from tqdm.auto import tqdm
from multiprocessing import Pool

import logging
logger = logging.getLogger(__name__)

class State(IntEnum):
    M = 0
    I = 1
    D = 2

class Transition(IntEnum):
    M2M = 0
    M2I = 1
    M2D = 2
    I2M = 3
    I2I = 4
    D2M = 5
    D2D = 6

class nt_index(IntEnum):
    A = 0
    T = 1
    G = 2
    C = 3
    PAD = 4
    SOS = 5
    EOS = 6
    U = 1

def one_hot_index(seq):
    return [int(nt_index[char]) for char in seq]

def one_hot_encode(nucleotide, padding=0):
    """入力された文字列に対してOne-hotなnp形式を返す
    """
    # パディングの大きさを指定する
    arr = np.vstack((np.eye(4), np.ones(4)[None, :]*0.25))
    return arr[one_hot_index("N"*padding + nucleotide + "N"*padding)].T

class SNV(IntEnum):
    Mutation = 0
    Insertion = 1
    Deletion = 2

class SequenceGenerator():
    def __init__(self, num_motifs=1, motif_length=10, motifs=None,
                 target_length=20, fix_random_region_length=True, error_rate=0, generate_motifs=True, middle_insert_range=[2, 6],
                 seed=0, add_primer=True, forward_primer="AAAAA", reverse_primer="GGGGG", one_side_proba=0.5, paired=False):
        np.random.seed(seed)

        if generate_motifs:
            self.motifs = ["".join(np.random.choice(
                list("ATGC"), motif_length)) for _ in range(num_motifs)]
        else:
            self.motifs = motifs

        self.error_indices = 1 + \
            np.argsort(np.random.random(size=motif_length-1))[:3]
        self.mut_idx, self.ins_idx, self.del_idx = self.error_indices

        logger.info(f"error rate is {error_rate*100:.1f}%")
        for idx, motif in enumerate(self.motifs):
            seq = [ch for ch in motif]
            mut = self.mutate(seq[self.mut_idx])
            if error_rate != 0:
                seq[self.mut_idx] = f"[{seq[self.mut_idx]}>{mut}]"
                seq[self.ins_idx] = f"[+]{seq[self.ins_idx]}"
                seq[self.del_idx] = f"{seq[self.del_idx].lower()}"
            seq = "".join(seq)
            logger.info(f"motif {idx} is {seq}")

        self.num_motifs = num_motifs
        self.error_rate = error_rate
        self.target_length = target_length
        self.forward_primer = forward_primer
        self.reverse_primer = reverse_primer
        self.add_primer = add_primer

        self.one_side_proba = one_side_proba
        self.middle_insert_range = middle_insert_range
        self.paired = paired

    def mutate(self, char):
        return "TGCA"["ATGC".index(char)]

    def sample_motif(self, n):
        motif_indices = np.random.randint(self.num_motifs, size=n)
        has_errors = np.random.random(size=n) < self.error_rate
        # mutation, insertion, deletion
        error_types = np.random.choice(SNV, size=n)
        sequences = []
        for motif_index, has_error, error_type in zip(motif_indices, has_errors, error_types):
            motif = self.motifs[motif_index]
            seq = [ch for ch in motif]
            if has_error:
                if error_type == SNV.Mutation:
                    seq[self.mut_idx] = self.mutate(seq[self.mut_idx])
                elif error_type == SNV.Insertion:
                    seq[self.ins_idx] = np.random.choice(
                        list("ATGC")) + seq[self.ins_idx]
                elif error_type == SNV.Deletion:
                    seq[self.del_idx] = ""
                else:
                    raise NotImplementedError
            seq = "".join(seq)
            sequences.append(seq)
        return sequences, motif_indices.tolist()

    def sample(self, n=1, with_indices=True):
        motifs, motif_indices = self.sample_motif(n)
        sequences = []
        paired_indices = []
        for seq in motifs:
            if self.paired:
                seq, idx = self.insert_in_the_middle(
                    seq, nrange=self.middle_insert_range, one_side_proba=self.one_side_proba)
                paired_indices += [idx]
            random_region = "".join(np.random.choice(
                list("ATGC"), size=self.target_length-len(seq)))
            l = np.random.randint(len(random_region))
            if self.add_primer:
                sequences.append(
                    self.forward_primer + random_region[:l] + seq + random_region[l:] + self.reverse_primer)
            else:
                sequences.append(random_region[:l] + seq + random_region[l:])

        if self.paired and with_indices:
            return sequences, motif_indices, paired_indices
        elif with_indices:
            return sequences, motif_indices
        return sequences

    def insert_in_the_middle(self, sequence, nrange=[2, 6], one_side_proba=0.5):
        n = np.random.randint(*nrange)
        if np.random.random() < one_side_proba:
            if np.random.choice(["l", "r"]) == "l":
                l_motif = sequence[:len(sequence)//2]
                r_motif = ""
                idx = 1
            else:
                l_motif = ""
                r_motif = sequence[len(sequence)//2:]
                idx = 2
        else:
            l_motif = sequence[:len(sequence)//2]
            r_motif = sequence[len(sequence)//2:]
            idx = 0
        return l_motif + "".join(np.random.choice(list("ATGC"), size=n)) + r_motif, idx

def get_reads_with_id_prefix(path, prefix_on, prefix_off):
    reads = []
    read = ""
    switch = False
    with path.open() as f:
        for line in f.readlines():
            if line[0] == prefix_off:
                switch = False
                reads.append(read)
                read = ""
            if switch:
                read = read + line.strip()
            if line[0] == prefix_on:
                switch = True
                read = ""
        # write last read line
        reads.append(read)
    return reads

def read_fasta(path):
    return get_reads_with_id_prefix(Path(path), ">", ">")

def read_fastq(path):
    return get_reads_with_id_prefix(Path(path), "@", "+")

class SingleRound:
    """pass path or raw_reads to make class of selex experiment per round.
    """

    def __init__(self, raw_reads: list = None, forward_adapter=None, reverse_adapter=None, name=None, tolerance=0, path: str = None, max_len=None):
        assert path is not None or raw_reads is not None, "either path or raw_reads has to be specified"
        if path:
            path = Path(path)
            if path .suffix == ".fastq":
                logger.info("reading fastq format sequence")
                raw_reads = read_fastq(path)
            elif path.suffix in {".fasta", ".fa"}:
                logger.info("reading fasta format sequence")
                raw_reads = read_fasta(path)
            else:
                logger.critical(
                    "please specify a file with fasta or fastq format")
                quit()

        self.raw_reads = raw_reads
        self.calc_target_length()
        self.max_len = max_len

        if forward_adapter is None or reverse_adapter is None:
            logger.info("adapter info not provided. estimating value")
            self.calc_experimental_settings()
        else:
            logger.info(
                f"sequence design : {forward_adapter}-[random]-{reverse_adapter}")
            self.set_adapters(forward_adapter, reverse_adapter,
                              self.max_len is not None)

        if name:
            self.name = name
        else:
            self.name = re.sub(r'[-\.\:]', "",
                               str(datetime.datetime.now())).replace(" ", "_")
        logger.info(f"experiment name : {self.name}")
        self.tolerance = tolerance

    def get_adapters(self):
        return self.forward_adapter, self.reverse_adapter

    def set_adapters(self, forward_adapter: str, reverse_adapter: str, set_max_len=False):
        self.forward_adapter = forward_adapter
        self.forward_adapter_length = len(forward_adapter)

        self.reverse_adapter = reverse_adapter
        self.reverse_adapter_length = len(reverse_adapter)

        self.random_region_length = self.target_length - \
            self.reverse_adapter_length - self.forward_adapter_length
        if set_max_len:
            self.random_region_length = self.max_len

    def calc_target_length(self):
        from collections import Counter, defaultdict
        self.read_counter = Counter(self.raw_reads)

        # calc most common length
        d = defaultdict(int)
        for key, value in self.read_counter.items():
            d[len(key)] += value
        self.target_length = sorted(d.items(), key=lambda x: -x[1])[0][0]

    def calc_experimental_settings(self):
        """calculate sequence adapters in a heuristic way
        """

        # fwd
        max_count = None
        est_adapter = ""
        for i in range(1, self.target_length):
            d = defaultdict(int)
            for seq, count in self.read_counter.most_common():
                if len(seq) < i or len(d) > 100 and seq[:i] not in d.keys():
                    continue
                d[seq[:i]] += count
            top_seq, top_count = sorted(d.items(), key=lambda x: -x[1])[0]
            if max_count is not None and top_count < max_count * 0.5:  # heuristics
                logger.info(
                    f"estimated forward adapter len is {i-1} : {est_adapter}")
                break
            max_count = sorted(d.items(), key=lambda x: -x[1])[0][1]
            if max_count < sum(self.read_counter.values()) * 0.5:
                logger.info(
                    f"no match found.")
                break
            est_adapter = top_seq
        fwd_len = i - 1
        fwd_adapter = est_adapter

        # rev
        max_count = None
        est_adapter = ""
        for i in range(1, self.target_length):
            d = defaultdict(int)
            for seq, count in self.read_counter.most_common():
                if len(seq) < i or len(d) > 100 and seq[-i:] not in d.keys():
                    continue
                d[seq[-i:]] += count
            top_seq, top_count = sorted(d.items(), key=lambda x: -x[1])[0]
            if max_count is not None and top_count < max_count * 0.5:  # heuristics
                logger.info(
                    f"estimated reverse adapter len is {i-1} : {est_adapter}")
                break
            max_count = sorted(d.items(), key=lambda x: -x[1])[0][1]
            if max_count < sum(self.read_counter.values()) * 0.5:
                logger.info(
                    f"no match found.")
                break
            est_adapter = top_seq
        rev_len = i - 1
        rev_adapter = est_adapter

        rand_len = self.target_length - rev_len - fwd_len

        logger.info(
            f"filtering with : {fwd_adapter}({fwd_len}N)-{rand_len}N-{rev_adapter}({rev_len}N)")

        # write estimated experimental settings
        self.set_adapters(fwd_adapter, rev_adapter, self.max_len is not None)

    def get_sequences_and_count(self):
        c = Counter(self.raw_reads)
        return c.most_common()

    def get_filter_passed_sequences_and_count(self, random_only=False):
        if random_only:
            return {self.cut_adapters(key): value for key, value in self.get_sequences_and_count()}
        else:
            c = Counter(self.get_filter_passed_sequences())
            return c.most_common()

    def filter_function(self, read):
        has_forward = read[: self.forward_adapter_length] == self.forward_adapter \
            or self.forward_adapter_length == 0
        has_reverse = read[-self.reverse_adapter_length:] == self.reverse_adapter \
            or self.reverse_adapter_length == 0
        match_random_region_len = abs(
            len(read) - self.target_length) <= self.tolerance
        return has_forward and has_reverse and match_random_region_len

    def get_filter_passed_sequences(self, random_only=False):
        self.filter_passed = list(filter(self.filter_function, self.raw_reads))
        if random_only:
            return [self.cut_adapters(read) for read in self.filter_passed]
        return self.filter_passed

    def cut_adapters(self, seq):
        if self.reverse_adapter_length == 0:
            ret = seq[self.forward_adapter_length:]
        else:
            ret = seq[self.forward_adapter_length: -
                      self.reverse_adapter_length]
        if self.max_len is not None:
            return ret[len(ret) // 2 - self.max_len // 2: len(ret) // 2 - self.max_len // 2 + self.max_len]
        else:
            return ret

    def __str__(self):
        return f"experiment of {len(self.raw_reads)} raw reads"

    def get_dataloader(self, min_count=1, test_size=0.1, batch_size=512, shuffle=True, use_cuda=True):
        from sklearn.model_selection import train_test_split
        from torch.utils.data import DataLoader

        self.min_count = min_count
        kwargs = {'num_workers': 1, 'pin_memory': True} if (
            use_cuda and torch.cuda.is_available()) else {}
        # load RAPT1-4R and filter reads to count>1, then make it to one hot encoded tensor
        c = self.get_filter_passed_sequences(random_only=True)
        sequences = list(
            filter(lambda seq_count: seq_count[1] >= min_count, Counter(c).most_common()))
        seq, _ = zip(*sequences)

        train_test = np.array(list(map(one_hot_index, seq)))
        logger.info(f"# of sequences -> {len(train_test)}")
        train_data, test_data = train_test_split(
            train_test, test_size=test_size, shuffle=shuffle)
        train_data = torch.from_numpy(train_data).long()
        test_data = torch.from_numpy(test_data).long()
        train_loader = DataLoader(
            train_data, batch_size=batch_size, shuffle=True,  **kwargs)
        test_loader = DataLoader(
            test_data,  batch_size=batch_size, shuffle=False, **kwargs)
        return train_loader, test_loader

class Dataset(torch.utils.data.Dataset):
    def __init__(self, data, transform=None):
        self.transform = transform
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        out_data = self.data[index]
        return out_data

def local_alignment(s1, s2, print_result=False, global_alignment=False):
    GAP_COST = -1
    MATCH_COST = +1
    MISMATCH_COST = -1

    # [[(0,0) for i in range(len(s2)+1)] for j in range(len(s1)+1)]
    dp = np.zeros((len(s1) + 1, len(s2) + 1, 2), dtype=np.int)

    def s(c1, c2):
        return MATCH_COST if c1 == c2 else MISMATCH_COST

    def idx_max(*args):
        argmax_i = np.argmax(args)
        return argmax_i, args[argmax_i]

    def concat(ret, a, b):
        return a + ret[0], b + ret[1]

    OBJ = 1
    POINTER_INDEX_TO_COND = ["MATCH", "S1_GAP", "S2_GAP", "END"]

    # DP
    for i1, c1 in enumerate(s1):
        for i2, c2 in enumerate(s2):
            l = [dp[i1][i2][OBJ] + s(c1, c2),
                 dp[i1][i2 + 1][OBJ] + GAP_COST,
                 dp[i1 + 1][i2][OBJ] + GAP_COST]
            if not global_alignment:
                l += [0]
            dp[i1 + 1][i2 + 1] = idx_max(*l)

    # local_alignment_traceback
    # logger.info(list(zip(*np.where(dp[:,:,OBJ]==dp[:,:,OBJ].max()))))
    result = []
    if global_alignment:
        traceback_starts = [(len(s1), len(s2))]
    else:
        traceback_starts = list(
            zip(*np.where(dp[:, :, OBJ] == dp[:, :, OBJ].max())))
    for i1, i2 in traceback_starts:
        traceback_pointer = dp[i1, i2, 0]
        ret = ["", ""]
        while True:
            if print_result:
                logger.info(i1, i2, POINTER_INDEX_TO_COND[traceback_pointer])
            if POINTER_INDEX_TO_COND[traceback_pointer] == "MATCH":
                i1 -= 1
                i2 -= 1
                if i1 < 0 or i2 < 0:
                    break
                if s1[i1] != s2[i2]:
                    ret = concat(ret, s1[i1].lower(), s2[i2].lower())
                else:
                    ret = concat(ret, s1[i1], s2[i2])
                traceback_pointer = dp[i1, i2, 0]
            elif POINTER_INDEX_TO_COND[traceback_pointer] == "S1_GAP":
                i1 -= 1
                if i1 < 0 or i2 < 0:
                    break
                ret = concat(ret, s1[i1], "-")
                traceback_pointer = dp[i1, i2, 0]
            elif POINTER_INDEX_TO_COND[traceback_pointer] == "S2_GAP":
                i2 -= 1
                if i1 < 0 or i2 < 0:
                    break
                ret = concat(ret, "-", s2[i2])
                traceback_pointer = dp[i1, i2, 0]
            else:
                break
        if print_result:
            logger.info(ret, dp[:, :, OBJ].max())
        result.append(ret)
    return dp[:, :, OBJ].max(), result

# from https://rosettacode.org/wiki/Levenshtein_distance#Memoized_recursive_version_2
@lru_cache(maxsize=2**26)
def edit_distance(s, t):
    if not s:
        return len(t)
    if not t:
        return len(s)
    if s[0] == t[0]:
        return edit_distance(s[1:], t[1:])
    l1 = edit_distance(s, t[1:])
    l2 = edit_distance(s[1:], t)
    l3 = edit_distance(s[1:], t[1:])
    return 1 + min(l1, l2, l3)

def get_complement_sequence(seq):
    return "".join(reversed([list("ATGC")["TACG".index(char)] for char in seq]))

class ProfileHMMSampler():
    def __init__(self, transition_proba, emission_proba, proba_is_log=False):
        self.e = emission_proba
        self.a = transition_proba
        if proba_is_log:
            self.e = np.exp(self.e)
            self.a = np.exp(self.a)
        self.e = self.e / np.sum(self.e, axis=1)[:, None]

    def sample(self, sequence_only=False, debug=False):
        idx, state = (0, State.M)
        states = [(idx, state)]
        seq = ""
        while True:
            if state == State.M:
                p = self.a[idx][np.array([
                    Transition.M2M.value,
                    Transition.M2I.value,
                    Transition.M2D.value])]
            elif state == State.I:
                p = np.stack([
                    self.a[idx][Transition.I2M.value],
                    self.a[idx][Transition.I2I.value],
                    0])
            elif state == State.D:
                p = np.stack([
                    self.a[idx][Transition.D2M.value],
                    0,
                    self.a[idx][Transition.D2D.value]])
            else:
                logger.info("something wrong")

            state = np.random.choice([State.M, State.I, State.D], p=p/sum(p))
            if state != State.I:
                idx += 1
            states.append((idx, state))
            if idx == self.a.shape[0]:
                break

            if state == State.M:
                # logger.info("{:.2f}, {:.2f}, {:.2f}, {:.2f}".format(*self.e[idx-1]))

                seq += np.random.choice(list("ATGC"), p=self.e[idx-1])
                if debug:
                    logger.info(idx, state, self.e[idx-1], seq[-1])
            elif state == State.I:
                seq += np.random.choice(list("atgc"))
            else:
                seq += "_"
        if not sequence_only:
            return states, seq
        else:
            return seq

    def most_probable(self, sequence_only=False):
        idx, state = (0, State.M)
        states = [(idx, state)]
        seq = ""
        while True:
            if state == State.M:
                p = self.a[idx][np.array([
                    Transition.M2M.value,
                    Transition.M2I.value,
                    Transition.M2D.value])]
            elif state == State.I:
                p = [
                    self.a[idx][Transition.I2M.value],
                    0,
                    0]
            elif state == State.D:
                p = [
                    self.a[idx][Transition.D2M.value],
                    0,
                    self.a[idx][Transition.D2D.value]]
            else:
                logger.info("something wrong")
            p[np.argmax(p)] += 1000000
            state = np.random.choice([State.M, State.I, State.D], p=p/sum(p))
            if state != State.I:
                idx += 1
            states.append((idx, state))

            if idx == self.a.shape[0]:
                break

            if state == State.M:
                # logger.info("{:.2f}, {:.2f}, {:.2f}, {:.2f}".format(*self.e[idx-1]))
                p = np.copy(self.e[idx-1])
                p[np.argmax(p)] += 100000
                seq += np.random.choice(list("ATGC"), p=p/sum(p))
            elif state == State.I:
                seq += "N"
            else:
                seq += "_"
        if not sequence_only:
            return states, seq
        else:
            return seq

    def calc_seq_proba(self, seq: str):
        one_hot_seq = torch.tensor(one_hot_index(seq))
        model_len = self.e.shape[0]
        random_len = len(seq)

        e = np.log(self.e)
        a = np.log(self.a)

        F = torch.ones((3, model_len + 2, random_len + 1)) * (-100)

        # init
        F[0, 0, 0] = 0

        for i in range(random_len + 1):
            for j in range(model_len + 1):
                # State M
                if j*i != 0:
                    F[State.M, j, i] = e[j - 1][one_hot_seq[i - 1]] + \
                        torch.logsumexp(torch.stack((
                            a[j - 1, Transition.M2M] +
                            F[State.M, j - 1, i - 1],
                            a[j - 1, Transition.I2M] +
                            F[State.I, j - 1, i - 1],
                            a[j - 1, Transition.D2M] + F[State.D, j - 1, i - 1])), dim=0)

                # State I
                if i != 0:
                    F[State.I, j, i] = - 1.3863 + \
                        torch.logsumexp(torch.stack((
                            a[j, Transition.M2I] + F[State.M, j, i-1],
                            a[j, Transition.I2I] + F[State.I, j, i-1]
                        )), dim=0)

                # State D
                if j != 0:
                    F[State.D, j, i] = \
                        torch.logsumexp(torch.stack((
                            a[j - 1, Transition.M2D] + F[State.M, j - 1, i],
                            a[j - 1, Transition.D2D] + F[State.D, j - 1, i]
                        )), dim=0)

        F[State.M, model_len+1, random_len] = \
            torch.logsumexp(torch.stack((
                a[model_len, Transition.M2M] +
                F[State.M, model_len, random_len],
                a[model_len, Transition.I2M] +
                F[State.I, model_len, random_len],
                a[model_len, Transition.D2M] +
                F[State.D, model_len, random_len]
            )), dim=0)

        return F[State.M, model_len+1, random_len]

class Result():
    """実験結果の保存のためのクラス"""
    from raptgen.visualization import provide_ax

    def __init__(self,
                 model,
                 path_to_selex: str = None,
                 experiment: SingleRound = None,
                 path_to_result_csv: str = None,
                 path_to_model: str = None,
                 lazy_mu_eval=False,
                 path_to_save_results=None,
                 evaluated_X=None,
                 evaluated_y=None,
                 load_if_exists=False,
                 min_count=1
                 ):

        if experiment is None:
            self.experiment = SingleRound(path=Path(path_to_selex))
        else:
            self.experiment = experiment

        if path_to_result_csv:
            self.result_df = pd.read_csv(Path(path_to_result_csv))
        else:
            logger.info("skip loading training result")

        if path_to_model:
            model.load_state_dict(torch.load(path_to_model))
        else:
            logger.info("skip loading model parameters")

        self.model = model

        if path_to_save_results is None:
            self.path_to_save_results = Path(
                "result_" + self.get_result_hash())
        else:
            self.path_to_save_results = Path(path_to_save_results)
        if load_if_exists:
            self.path_to_save_results.mkdir(parents=True, exist_ok=True)
        else:
            self.path_to_save_results.mkdir(parents=True)

        self.is_phmm = "PHMM" in str(model.__class__)
        self.min_count = min_count
        if not lazy_mu_eval:
            logger.info("evaluating mu")
            self.get_mean_vectors_from_experiment()

        self.evaluated_X = evaluated_X
        self.evaluated_y = evaluated_y

    def get_mean_vectors_from_experiment(self, get_raw_seq=False, force=False):
        if not hasattr(self, "mus") or not hasattr(self, "seqs") or force:
            loaders = self.experiment.get_dataloader(
                shuffle=False, min_count=self.min_count)
            with torch.no_grad():
                self.model.eval()
                mus = []
                seqs = []
                for loader in loaders:
                    for data in loader:
                        _, mu, logvar = self.model(data, deterministic=True)
                        mus += [*mu.detach().numpy()]
                        for datum in data:
                            seq = "".join(np.array(list("ATGC"))
                                          [datum.numpy()])
                            seqs.append(seq)

            self.mus = np.stack(mus)
            self.seqs = seqs

        if get_raw_seq:
            return self.mus, self.seqs
        return self.mus

    def get_result_hash(self):
        # modelのパラメタが一意に決まっていてhashしやすいのでこれを利用する
        import hashlib
        b = "".join(["{:.2f}".format(i.flatten()[0].cpu().detach().numpy())
                     for i in self.model.parameters()])
        self.hash = hashlib.sha1(b.encode()).hexdigest()[:10]
        logger.info(f"hash : {self.hash}")
        return self.hash

    def calc_gmm(self, dim=10, calc_times=100, force=False):
        from sklearn.mixture import GaussianMixture
        logger.info("calculating gmm centers")
        X = self.mus
        gmm_path = self.path_to_save_results/"gmm.pkl"
        if gmm_path.exists():
            logger.info(f"loading {gmm_path}")
            with gmm_path.open("rb") as f:
                best_gmm = pickle.load(f)
                best_aic = best_gmm.aic(X)
        else:
            best_aic = np.inf
            pbar = tqdm(range(calc_times))
            for i in pbar:
                gmm = GaussianMixture(dim, covariance_type="full").fit(X)
                if gmm.aic(X) < best_aic:
                    best_aic = gmm.aic(X)
                    best_gmm = gmm
                pbar.set_description(
                    "[" + "⠸⠴⠦⠇⠋⠙"[i % 6] + "]" + f"{best_aic:.2f}")

            with gmm_path.open("wb") as f:
                pickle.dump(best_gmm, f)

        logger.info(f"best aic : {best_aic}")
        self.gmm = best_gmm
        self.aic = best_aic
        self.gmm_classes = best_gmm.predict(X)
        self.gmm_centers = best_gmm.means_

    def embed_sequences(self, sequences):
        """
        docstring
        """
        if type(sequences) == str:
            sequences = [sequences]

        # https://discuss.pytorch.org/t/how-to-check-if-model-is-on-cuda/180
        model_device = next(self.model.parameters()).device
        with torch.no_grad():
            self.model.eval()
            mus = []
            for sequence in sequences:
                recon, mu, logvar = self.model(
                    torch.Tensor(
                        [one_hot_index(sequence)],
                        device=model_device).long())
                mus += [mu]
        return torch.cat(mus)

    @provide_ax
    def plot_gmm(self, ax, fig=None, save=True, no_colors=False, no_gmm_centers=False):
        if not hasattr(self, "gmm"):
            logger.info("calculating gmm")
            self.calc_gmm()
        if len(self.gmm_classes) > 1000:
            if no_colors:
                ax.scatter(*self.mus[:1000].T, s=2, c="silver")
            else:
                ax.scatter(*self.mus[:1000].T, c=self.gmm_classes[:1000], s=2)
        else:
            if no_colors:
                ax.scatter(*self.mus.T, s=2, c="silver")
            else:
                ax.scatter(*self.mus.T, c=self.gmm_classes, s=2)
        if not no_gmm_centers:
            ax.scatter(*self.gmm_centers.T, c="r", marker="*",
                       s=10, zorder=50, label="gmm center")
            for i, (x, y) in enumerate(self.gmm_centers):
                ax.text(x, y, f" {i}", ha="left", va="center", c="r", bbox=dict(
                    facecolor='white', alpha=0.5), zorder=40)

        ax.axis("square")
        if save:
            fig.savefig(self.path_to_save_results/'gmm.png')
        return ax

    @provide_ax
    def plot_means(self, ax, fig=None, with_count=False, save=True, meshgrid=True):
        if self.mus is None:
            logger.info("running estimation")
            self.get_mean_vectors_from_experiment()

        if with_count:
            c = Counter(
                self.experiment.get_filter_passed_sequences(random_only=True))
            counts = [np.log2(c[seq]) for seq in self.seqs]
            XY = self.mus[np.argsort(counts)]
            if meshgrid:
                XY_ = []
                for xy, cnt in zip(XY, [c[seq] for seq in self.seqs]):
                    XY_ += [xy]*int(np.log2(cnt))
                heatmap, xedges, yedges = np.histogram2d(
                    * np.stack(XY_).T, bins=100, range=((-3, 3), (-3, 3)))
            else:
                cs = ax.scatter(*XY.T, s=2, c=sorted(counts), vmin=0)
                cbar = fig.colorbar(cs)
                cbar.ax.set_title("log2 counts")
                legend = ax.legend(*cs.legend_elements(num=6),
                                   loc="lower right", title="log2 count")
                ax.add_artist(legend)

        else:
            ax.scatter(*self.mus.T, s=2)
        ax.axis("square")

        if save:
            fig.savefig(self.path_to_save_results/'means.png')
        return ax

    @provide_ax
    def plot_bo(self, ax, fig=None, n_grid=101, save=True, with_index=True, plot_range=(-2, 2)):
        from mpl_toolkits.axes_grid1 import make_axes_locatable
        if not hasattr(self, "next_locations"):
            self.get_bo_result()

        self.bo_grid_x, self.bo_grid_y = np.meshgrid(
            np.linspace(*plot_range, n_grid),
            np.linspace(*plot_range, n_grid))
        self.bo_X = np.stack(map(np.ravel, (self.bo_grid_x, self.bo_grid_y))).T

        mu, sigma = self.bo.model.predict(self.bo_X)
        self.bo_mu = mu
        self.bo_sigma = sigma

        cont = ax.contour(self.bo_grid_x, self.bo_grid_y, -
                          (self.bo_mu-self.bo_sigma).reshape(n_grid, n_grid))
        cont.clabel(fmt='%1.1f', fontsize=8)
        divider = make_axes_locatable(ax)
        cax = divider.append_axes('right', size='5%', pad=0.05)

        ax.scatter(*self.evaluated_X.T,
                   c=self.evaluated_y[:, 0], cmap="bwr_r", ec="grey", lw=0.5, zorder=40, label="evaluated")
        if with_index:
            for i, (x, y) in enumerate(self.evaluated_X):
                ax.text(x, y, " "+str(i), color="blue", va="center")

        ax.scatter(*self.next_locations.T, marker="*",
                   color="k", zorder=50, label="bo proposed")
        if with_index:
            for i, (x, y) in enumerate(self.next_locations):
                ax.text(x, y, f" {i}", color="k", zorder=51, va="center")
        fig.colorbar(cont, cax=cax)

        return ax

    def get_bo_result(self, n=10, domain=(-2, 2), force_rerun=False):
        if hasattr(self, "next_locations") and not force_rerun:
            return self.next_locations
        import GPyOpt
        bo_path = self.path_to_save_results/"bo.pkl"
        if bo_path.exists() and not force_rerun:
            logger.info(f"loading {bo_path}")
            with bo_path.open("rb") as f:
                self.bo = pickle.load(f)
        else:
            logger.info("calculating bo")
            assert self.evaluated_X is not None and self.evaluated_y is not None,\
                "(N, d) array: `evaluated_X` and (N, 1) array: `evaluated_y` should be set"
            self.domain = domain
            self.constraints = [{"name": f"var_{i+1}",
                                 "type": "continuous",
                                 "domain": self.domain}
                                for i in range(self.evaluated_X.shape[1])]

            self.bo = GPyOpt.methods.BayesianOptimization(
                None,
                domain=self.constraints,
                model_type='GP',
                acquisition_type='LCB',
                evaluator_type="local_penalization",
                batch_size=n,
                X=self.evaluated_X,
                Y=self.evaluated_y)
            with bo_path.open("wb") as f:
                pickle.dump(self.bo, f)

        self.next_locations = self.bo.suggest_next_locations()
        return self.next_locations

    def _points_to_score(self, points, eval_max=256):
        a, e_m = self.model.decoder(points)
        a = a.detach().numpy()
        e_m = e_m.detach().numpy()

        logger.info(
            f"calculating most probable sequences up to {eval_max} candidates")
        pbar = tqdm(range(len(points)))
        scores = []
        for j in pbar:
            sampler = ProfileHMMSampler(a[j], e_m[j], proba_is_log=True)
            seq_pattern = sampler.most_probable()[1].replace(
                "_", "").replace("N", "*")
            products = product(*[list("ATGC")
                                 for _ in range(seq_pattern.count("*"))])

            rets = []
            for nt_set in products:
                ret = ""
                for part, nt in zip(seq_pattern.split("*"), list(nt_set)+[""]):
                    ret += part+nt
                rets += [ret]
            if len(rets) > eval_max:
                rets = [rets[idx] for idx in np.argsort(
                    np.random.randn(len(rets)))[:eval_max]]
            with Pool() as p:
                probas = p.map(sampler.calc_seq_proba, rets)

            most_probable_seq, min_value = sorted(
                list(zip(rets, probas)), key=lambda x: x[1])[0]
            min_value = min_value.item()
            scores += [(seq_pattern, most_probable_seq, min_value)]
        self.scores = scores
        return scores

    def _save_scores(self, scores, model_type, filename, id_header="", force=False, loc=None):
        from datetime import datetime
        if filename is not None:
            assert not ((not force) and (self.path_to_save_results /
                                         filename).exists()), "file exists. to override, try: 'force=True'"
            with open(self.path_to_save_results / filename, "w") as f:
                f.write("id,method,max_model,max_seq,log_proba")
                if loc is not None:
                    f.write(",pos_x,pos_y\n")
                else:
                    f.write("\n")

                for i, (seq_pattern, most_probable_seq, min_value) in enumerate(scores):
                    id_str = datetime.now().strftime(
                        "%-y%m%d_") + id_header + f"_{i}"
                    f.write(
                        f"{id_str},{model_type},{seq_pattern},{most_probable_seq},{min_value}")
                    if loc is not None:
                        f.write(f",{loc[i][0]},{loc[i][1]}\n")
                    else:
                        f.write("\n")
            logger.info(f"saved to {self.path_to_save_results}/{filename}")

    def get_gmm_probable_sequences(self, filename=None):
        if hasattr(self, "scores"):
            return [most_probable for seq_pattern, most_probable, min_value in self.scores]
        if not hasattr(self, "gmm_centers"):
            self.calc_gmm()

        scores = self._points_to_score(torch.Tensor(self.gmm_centers))

        self._save_scores(scores, "GMM", filename)

        return [most_probable for seq_pattern, most_probable, min_value in scores]

    def plot_training_result(self, nwarmup=100, save=True, fig=None, axes=None):
        from raptgen.visualization import get_ax
        if axes is not None and fig is not None:
            ax, ay = axes
        fig, (ax, ay) = get_ax(row_col=(2, 1), return_fig=True)

        # for mean plot
        length = 20
        v = np.sin(np.arange(length)/(length-1)*np.pi)
        v /= sum(v)

        i = 0
        for arr_name in ["test_loss", "test_recon", "train_loss"]:
            cmap = plt.get_cmap("Paired")
            test_losses = np.array(self.result_df[arr_name])
            epochs = np.arange(len(test_losses))
            conv_indices, conv_values = np.array(
                epochs[:-length+1])+length//2, np.convolve(test_losses, v, mode='valid')
            ax.plot(test_losses, c=cmap(i))

            ax.plot(conv_indices, conv_values, c=cmap(i+1), label=arr_name)
            ax.plot(conv_indices[-1], conv_values[-1], marker='.', c=cmap(i+1))
            ax.text(np.argmin(test_losses) + 1, np.min(test_losses),
                    f"←{np.min(test_losses):.2f}", ha="left", va="top", c=cmap(i+1), rotation=-45, fontsize=8,
                    bbox=dict(ec=(1, 1, 1, 1), facecolor="w", alpha=1, pad=0))
            i += 2

        for arr_name in ["test_kld"]:
            cmap = plt.get_cmap("Paired")
            test_losses = np.array(self.result_df[arr_name])
            epochs = np.arange(len(test_losses))
            conv_indices, conv_values = np.array(
                epochs[:-length+1])+length//2, np.convolve(test_losses, v, mode='valid')
            ay.plot(test_losses, c=cmap(i))

            ay.plot(conv_indices, conv_values, c=cmap(i+1), label=arr_name)
            ay.plot(conv_indices[-1], conv_values[-1], marker='.', c=cmap(i+1))

            i += 2

        min_test_loss = min(self.result_df.test_loss)
        min_test_recon = min(self.result_df.test_recon)
        dloss = abs(min_test_loss-min_test_recon)
        ax.set_ylim(min_test_recon-dloss*0.6, min_test_loss+dloss*0.6)
        ay.set_ylim(-dloss*0.1, dloss*2.1)

        ax.plot((nwarmup, nwarmup), (0, 100), "--",
                c="gray", label="profile hmm warmup")
        ay.plot((nwarmup, nwarmup), (0, 100), "--",
                c="gray", label="profile hmm warmup")
        ay.set_xlabel("epochs")
        ax.set_ylabel("loss")
        ay.set_ylabel("loss")

        ax.legend()
        ay.legend()
        ax.set_title("training result")
        if save:
            fig.savefig(self.path_to_save_results/'training.png')

class Experiments():
    """whole selex experiment"""

    def __init__(self, read_paths: list, has_same_adapters=True, k=3):
        adapters = None
        self.rounds = []
        for read_path in sorted(read_paths):
            logger.info(f"reading ... {read_path}")
            path = Path(read_path)
            if adapters is not None and has_same_adapters:
                fwd, rev = adapters
                single_round = SingleRound(path=path, name=path.stem,
                                           forward_adapter=fwd, reverse_adapter=rev)
            else:
                single_round = SingleRound(path=path, name=path.stem)
                adapters = single_round.get_adapters()
            self.rounds += [single_round]
        self.k = k
        self.kmer_list = sorted(list("".join(l)
                                     for l in product(list("ATGC"), repeat=k)))

    def has_exact_match(self, sequences):
        if type(sequences) is str:
            sequences = [sequences]
        results = dict()
        for sequence in sequences:
            appeared = set()
            for experiment in self.rounds:
                if sequence in set(experiment.get_filter_passed_sequences(random_only=True)):
                    appeared |= {experiment.name}
            results[sequence] = appeared
        return results

    def kmer(self, seq):
        return [seq[i: i+self.k] for i in range(len(seq)-self.k+1)]

    def kmer_count(self, seq, to_list=False):
        from collections import Counter
        c = Counter(self.kmer(seq))
        if to_list:
            return [c[kmer] for kmer in self.kmer_list]
        return c

    def save_frequencies(self, save_path, sequences, min_count=2, idx_header=""):
        from functools import partial

        if type(sequences) is str:
            sequences = [sequences]
        pbar = tqdm(total=len(sequences))

        whole_read_counts = Counter(
            [seq for experiment in self.rounds for seq in experiment.get_filter_passed_sequences(random_only=True)])
        whole_reads = [x[0] for x in filter(
            lambda x:x[1] >= min_count, whole_read_counts.most_common())]
        pbar.set_description("calculating kmer distribution")
        if not hasattr(self, "whole_reads_kmer"):
            self.whole_reads_kmer = np.stack(
                [np.array(self.kmer_count(seq, to_list=True)) for seq in whole_reads])

        with open(save_path, "w") as f:
            f.write(
                "idx\tgenerate_seq\texact_match_in\tnearest_selex_seq\tscore\tselex_aligned\tgenerate_aligned\t")
            f.write(
                "\t".join([experiment.name for experiment in self.rounds])+"\n")

            # for sequence in query
            for idx, most_probable_seq in enumerate(sequences):
                pbar.set_description("taking alignments")
                func = partial(local_alignment, s2=most_probable_seq)
                with Pool() as p:
                    alignments = p.map(func, whole_reads)
                arr = np.array(list(zip(*alignments))[0])

                matches = self.has_exact_match(most_probable_seq)[
                    most_probable_seq]

                pbar.set_description("writing to file")
                for jdx in np.argwhere(arr == max(arr)).flatten():
                    score = alignments[jdx][0]
                    opt_seq_align, check_seq_align = alignments[jdx][1][0]
                    f.write(
                        f"{idx_header}{idx}\t{most_probable_seq}\t{matches}\t{whole_reads[jdx]}\t{score}\t{opt_seq_align}\t{check_seq_align}")

                    for experiment in self.rounds:
                        c = experiment.get_filter_passed_sequences_and_count(
                            random_only=True)
                        if whole_reads[jdx] in c.keys():
                            ratio = f"\t{c[whole_reads[jdx]] / len(experiment.get_filter_passed_sequences()):.6f}"
                            f.write(ratio)
                        else:
                            f.write("\t0.000000")
                    f.write("\n")
                pbar.update(1)
        return pd.read_table(save_path)


# -*- coding: utf-8 -*-
import logging
import warnings
import numpy as np
from scipy.interpolate import RectBivariateSpline
import xarray as xr
import dask
import rasterio
import rasterio.features
import rioxarray
from scipy.interpolate import interp1d
from shapely.geometry import Polygon, box
import shapely
from .utils import timing, haversine, map_blocks_coords, bbox_coords, BlockingActorProxy, merge_yaml, get_glob, \
    to_lon180
from numpy import asarray
from affine import Affine
from .sentinel1_meta import Sentinel1Meta
from .ipython_backends import repr_mimebundle
import yaml

logger = logging.getLogger('xsar.sentinel1_dataset')
logger.addHandler(logging.NullHandler())

# we know tiff as no geotransform : ignore warning
warnings.filterwarnings("ignore", category=rasterio.errors.NotGeoreferencedWarning)

# allow nan without warnings
# some dask warnings are still non filtered: https://github.com/dask/dask/issues/3245
np.errstate(invalid='ignore')

class Sentinel1Dataset:
    """
    Handle a SAFE subdataset.
    A dataset might contain several tiff files (multiples polarizations), but all tiff files must share the same footprint.

    The main attribute useful to the end-user is `self.dataset` (`xarray.Dataset` , with all variables parsed from xml and tiff files.)

    Parameters
    ----------
    dataset_id: str or Sentinel1Meta object

        if str, it can be a path, or a gdal dataset identifier like `'SENTINEL1_DS:%s:WV_001' % filename`)

    resolution: dict, number or string, optional
        resampling dict like `{'atrack': 20, 'xtrack': 20}` where 20 is in pixels.

        if a number, dict will be constructed from `{'atrack': number, 'xtrack': number}`

        if str, it must end with 'm' (meters), like '100m'. dict will be computed from sensor pixel size.

    resampling: rasterio.enums.Resampling or str, optional

        Only used if `resolution` is not None.

        ` rasterio.enums.Resampling.rms` by default. `rasterio.enums.Resampling.nearest` (decimation) is fastest.

    luts: bool, optional

        if `True` return also luts as variables (ie `sigma0_lut`, `gamma0_lut`, etc...). False by default.

    chunks: dict, optional

        dict with keys ['pol','atrack','xtrack'] (dask chunks).

    dtypes: None or dict, optional

        Specify the data type for each variable.

    patch_variable: bool, optional

        activate or not variable pathching ( currently noise lut correction for IPF2.9X)

    See Also
    --------
    xsar.open_dataset
    """

    def __init__(self, dataset_id, resolution=None,
                 resampling=rasterio.enums.Resampling.rms,
                 luts=False, chunks={'atrack': 5000, 'xtrack': 5000},
                 dtypes=None, patch_variable=True):

        # miscellaneous attributes that are not know from xml files
        attrs_dict = {
            'pol': {
                'comment': 'ordered polarizations (copol, crosspol)'
            },
            'atrack': {
                'units': '1',
                'comment': 'azimuth direction, in pixels from full resolution tiff'
            },
            'xtrack': {
                'units': '1',
                'comment': 'cross track direction, in pixels from full resolution tiff'
            },
            'sigma0_raw': {
                'units': 'm2/m2'
            },
            'gamma0_raw': {
                'units': 'm2/m2'
            },
            'nesz': {
                'units': 'm2/m2',
                'comment': 'sigma0 noise'
            },
            'negz': {
                'units': 'm2/m2',
                'comment': 'beta0 noise'
            },
        }

        # default dtypes
        self._dtypes = {
            'latitude': 'f4',
            'longitude': 'f4',
            'incidence': 'f4',
            'elevation': 'f4',
            'altitude': 'f4',
            'ground_heading': 'f4',
            'nesz': None,
            'negz': None,
            'sigma0_raw': None,
            'gamma0_raw': None,
            'noise_lut': 'f4',
            'noise_lut_range': 'f4',
            'noise_lut_azi': 'f4',
            'sigma0_lut': 'f8',
            'gamma0_lut': 'f8',
            'azimuth_time': np.datetime64,
            'slant_range_time': None
        }
        if dtypes is not None:
            self._dtypes.update(dtypes)

        # default meta for map_blocks output.
        # as asarray is imported from numpy, it's a numpy array.
        # but if later we decide to import asarray from cupy, il will be a cupy.array (gpu)
        self._default_meta = asarray([], dtype='f8')

        self.s1meta = None
        """`xsar.Sentinel1Meta` object"""

        if not isinstance(dataset_id, Sentinel1Meta):
            self.s1meta = BlockingActorProxy(Sentinel1Meta, dataset_id)
            # check serializable
            # import pickle
            # s1meta = pickle.loads(pickle.dumps(self.s1meta))
            # assert isinstance(s1meta.coords2ll(100, 100),tuple)
        else:
            # we want self.s1meta to be a dask actor on a worker
            self.s1meta = BlockingActorProxy(Sentinel1Meta.from_dict, dataset_id.dict)
        del dataset_id

        if self.s1meta.multidataset:
            raise IndexError(
                """Can't open an multi-dataset. Use `xsar.Sentinel1Meta('%s').subdatasets` to show availables ones""" % self.s1meta.path
            )

        self._dataset = self._load_digital_number(resolution=resolution, resampling=resampling, chunks=chunks)
        self._dataset = xr.merge([xr.Dataset({'time': self._burst_azitime}), self._dataset])

        # dataset no-pol template for function evaluation on coordinates (*no* values used)
        # what's matter here is the shape of the image, not the values.
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", np.ComplexWarning)
            if self.s1meta._bursts['burst'].size != 0:
                # SLC TOPS, tune the high res grid because of bursts overlapping
                atrack_time = self._burst_azitime
                self._da_tmpl = xr.DataArray(
                    dask.array.empty_like(
                        np.empty((len(atrack_time), len(self._dataset.digital_number.xtrack))),
                        dtype=np.int8, name="empty_var_tmpl-%s" % dask.base.tokenize(self.s1meta.name)),
                    dims=('atrack', 'xtrack'),
                    coords={
                        'atrack': self._dataset.digital_number.atrack,
                        'xtrack': self._dataset.digital_number.xtrack,
                        'atrack_time': atrack_time.astype(float),
                    },
                )
            else:

                self._da_tmpl = xr.DataArray(
                    dask.array.empty_like(
                        self._dataset.digital_number.isel(pol=0).drop('pol'),
                        dtype=np.int8, name="empty_var_tmpl-%s" % dask.base.tokenize(self.s1meta.name)),
                    dims=('atrack', 'xtrack'),
                    coords={'atrack': self._dataset.digital_number.atrack,
                            'xtrack': self._dataset.digital_number.xtrack}
                )

        # FIXME possible memory leak
        # when calling a self.s1meta method, an ActorFuture is returned.
        # But this seems to break __del__ methods from both Sentinel1Meta and XmlParser
        # Is it a memory leak ?
        # see https://github.com/dask/distributed/issues/5610
        # tmp_f = self.s1meta.to_dict("all")
        # del tmp_f
        # return

        self._dataset.attrs.update(self.s1meta.to_dict("all"))

        # dict mapping for variables names to create by applying specified lut on digital_number
        self._map_var_lut = {
            'sigma0_raw': 'sigma0_lut',
            'gamma0_raw': 'gamma0_lut',
        }

        # dict mapping for lut names to file type (from self.files columns)
        self._map_lut_files = {
            'sigma0_lut': 'calibration',
            'gamma0_lut': 'calibration',
            'noise_lut_range': 'noise',
            'noise_lut_azi': 'noise',
        }

        # dict mapping specifying if the variable has 'pol' dimension
        self._vars_with_pol = {
            'sigma0_lut': True,
            'gamma0_lut': True,
            'noise_lut_range': True,
            'noise_lut_azi': True,
            'incidence': False,
            'elevation': False,
            'altitude': False,
            'azimuth_time': False,
            'slant_range_time': False,
            'longitude': False,
            'latitude': False
        }

        # variables not returned to the user (unless luts=True)
        self._hidden_vars = ['sigma0_lut', 'gamma0_lut', 'noise_lut', 'noise_lut_range', 'noise_lut_azi']
        # attribute to activate correction on variables, if available
        self._patch_variable = patch_variable

        self._luts = self._lazy_load_luts(self._map_lut_files.keys())

        # noise_lut is noise_lut_range * noise_lut_azi
        if 'noise_lut_range' in self._luts.keys() and 'noise_lut_azi' in self._luts.keys():
            self._luts = self._luts.assign(noise_lut=self._luts.noise_lut_range * self._luts.noise_lut_azi)
            self._luts.noise_lut.attrs['history'] = merge_yaml(
                [self._luts.noise_lut_range.attrs['history'] + self._luts.noise_lut_azi.attrs['history']],
                section='noise_lut'
            )

        self._rasterized_masks = self._load_rasterized_masks()

        ds_merge_list = [self._dataset, self._rasterized_masks, self._load_ground_heading(),  # lon_lat
                         self._luts.drop_vars(self._hidden_vars, errors='ignore')]

        if luts:
            ds_merge_list.append(self._luts[self._hidden_vars])
        attrs = self._dataset.attrs
        self._dataset = xr.merge(ds_merge_list)
        self._dataset.attrs = attrs

        for var_name, lut_name in self._map_var_lut.items():
            if lut_name in self._luts:
                # merge var_name into dataset (not denoised)
                self._dataset = self._dataset.merge(self._apply_calibration_lut(var_name))
                # merge noise equivalent for var_name (named 'ne%sz' % var_name[0)
                self._dataset = self._dataset.merge(self._get_noise(var_name))
            else:
                logger.debug("Skipping variable '%s' ('%s' lut is missing)" % (var_name, lut_name))

        self._dataset = self._dataset.merge(self._load_from_geoloc(['altitude', 'azimuth_time', 'slant_range_time',
                                                                    'incidence', 'elevation', 'longitude', 'latitude']))

        rasters = self._load_rasters_vars()
        if rasters is not None:
            self._dataset = xr.merge([self._dataset, rasters])

        self._dataset = self._dataset.merge(self._get_sensor_velocity())
        self._dataset = self._dataset.merge(self._range_ground_spacing())
        self._dataset = self._add_denoised(self._dataset)
        self._dataset.attrs = self._recompute_attrs()

        # set miscellaneous attrs
        for var, attrs in attrs_dict.items():
            try:
                self._dataset[var].attrs.update(attrs)
            except KeyError:
                pass

        self.sliced = False
        """True if dataset is a slice of original L1 dataset"""

        self.resampled = resolution is not None
        """True if dataset is not a sensor resolution"""

        # save original bbox
        self._bbox_coords_ori = self._bbox_coords

    def __del__(self):
        logger.debug('__del__')

    @property
    def dataset(self):
        """
        `xarray.Dataset` representation of this `xsar.Sentinel1Dataset` object.
        This property can be set with a new dataset, if the dataset was computed from the original dataset.
        """
        return self._dataset

    @dataset.setter
    def dataset(self, ds):
        if self.s1meta.name == ds.attrs['name']:
            # check if new ds has changed coordinates
            if not self.sliced:
                self.sliced = any(
                    [list(ds[d].values) != list(self._dataset[d].values) for d in ['atrack', 'xtrack', 'pol']])
            self._dataset = ds
            self._dataset.attrs = self._recompute_attrs()
        else:
            raise ValueError("dataset must be same kind as original one.")

    @dataset.deleter
    def dataset(self):
        logger.debug('deleter dataset')

    @property
    def _bbox_coords(self):
        """
        Dataset bounding box, in atrack/xtrack coordinates
        """
        bbox_ext = bbox_coords(self.dataset.atrack.values, self.dataset.xtrack.values)
        return bbox_ext

    @property
    def _bbox_ll(self):
        """Dataset bounding box, lon/lat"""
        return self.s1meta.coords2ll(*zip(*self._bbox_coords))

    @property
    def geometry(self):
        """
        geometry of this dataset, as a `shapely.geometry.Polygon` (lon/lat coordinates)
        """
        return Polygon(zip(*self._bbox_ll))

    @property
    def footprint(self):
        """alias for `xsar.Sentinel1Dataset.geometry`"""
        return self.geometry

    def ll2coords(self, *args):
        """
        Get `(atracks, xtracks)` from `(lon, lat)`,
        or convert a lon/lat shapely shapely object to atrack/xtrack coordinates.

        Parameters
        ----------
        *args: lon, lat or shapely object

            lon and lat might be iterables or scalars

        Returns
        -------
        tuple of np.array or tuple of float (atracks, xtracks) , or a shapely object

        Notes
        -----
        The difference with `xsar.Sentinel1Meta.ll2coords` is that coordinates are rounded to the nearest dataset coordinates.

        See Also
        --------
        xsar.Sentinel1Meta.ll2coords

        """
        if isinstance(args[0], shapely.geometry.base.BaseGeometry):
            return self.s1meta.ll2coords_shapely(args[0].intersection(self.geometry))

        atrack, xtrack = self.s1meta.ll2coords(*args)

        if hasattr(args[0], '__iter__'):
            scalar = False
        else:
            scalar = True

        tolerance = np.max([np.percentile(np.diff(self.dataset[c].values), 90) / 2 for c in ['atrack', 'xtrack']]) + 1
        try:
            # select the nearest valid pixel in ds
            ds_nearest = self.dataset.sel(atrack=atrack, xtrack=xtrack, method='nearest', tolerance=tolerance)
            if scalar:
                (atrack, xtrack) = (ds_nearest.atrack.values.item(), ds_nearest.xtrack.values.item())
            else:
                (atrack, xtrack) = (ds_nearest.atrack.values, ds_nearest.xtrack.values)
        except KeyError:
            # out of bounds, because of `tolerance` keyword
            (atrack, xtrack) = (atrack * np.nan, xtrack * np.nan)

        return atrack, xtrack

    def coords2ll(self, *args, **kwargs):
        """
         Alias for `xsar.Sentinel1Meta.coords2ll`

         See Also
         --------
         xsar.Sentinel1Meta.coords2ll
        """
        return self.s1meta.coords2ll(*args, **kwargs)

    @property
    def len_atrack_m(self):
        """atrack length, in meters"""
        bbox_ll = list(zip(*self._bbox_ll))
        len_m, _ = haversine(*bbox_ll[1], *bbox_ll[2])
        return len_m

    @property
    def len_xtrack_m(self):
        """xtrack length, in meters """
        bbox_ll = list(zip(*self._bbox_ll))
        len_m, _ = haversine(*bbox_ll[0], *bbox_ll[1])
        return len_m

    @property
    def pixel_atrack_m(self):
        """atrack pixel spacing, in meters (relative to dataset)"""
        return self.len_atrack_m / self.dataset.atrack.size

    @property
    def pixel_xtrack_m(self):
        """xtrack pixel spacing, in meters (relative to dataset)"""
        return self.len_xtrack_m / self.dataset.xtrack.size

    @property
    def coverage(self):
        """coverage string"""
        return "%dkm * %dkm (atrack * xtrack )" % (self.len_atrack_m / 1000, self.len_xtrack_m / 1000)

    @property
    def _regularly_spaced(self):
        return max(
            [np.unique(np.round(np.diff(self._dataset[dim].values), 1)).size for dim in ['atrack', 'xtrack']]) == 1

    def _recompute_attrs(self):
        if not self._regularly_spaced:
            warnings.warn(
                "Irregularly spaced dataset (probably multiple selection). Some attributes will be incorrect.")
        attrs = self._dataset.attrs
        attrs['pixel_xtrack_m'] = self.pixel_xtrack_m
        attrs['pixel_atrack_m'] = self.pixel_atrack_m
        attrs['coverage'] = self.coverage
        attrs['footprint'] = self.footprint
        return attrs

    def _patch_lut(self, lut):
        """
        patch proposed by MPC Sentinel-1 : https://jira-projects.cls.fr/browse/MPCS-2007 for noise vectors of WV SLC IPF2.9X products
        adjustement proposed by BAE are the same for HH and VV, and suppose to work for both old and new WV2 EAP
        they were estimated using WV image with very low NRCS (black images) and computing std(sigma0).
        Parameters
        ----------
        lut xarray.Dataset

        Returns
        -------
        lut xarray.Dataset
        """
        if self.s1meta.swath == 'WV':
            if lut.name in ['noise_lut_azi'] and self.s1meta.ipf in [2.9, 2.91] and \
                    self.s1meta.platform in ['SENTINEL-1A', 'SENTINEL-1B']:
                noise_calibration_cst_pp1 = {
                    'SENTINEL-1A':
                        {'WV1': -38.13,
                         'WV2': -36.84
                         },
                    'SENTINEL-1B':
                        {'WV1': -39.30,
                         'WV2': -37.44,
                         }
                }
                cst_db = noise_calibration_cst_pp1[self.s1meta.platform][self.s1meta.image['swath_subswath']]
                cst_lin = 10 ** (cst_db / 10)
                lut = lut * cst_lin
                lut.attrs['comment'] = 'patch on the noise_lut_azi : %s dB' % cst_db
        return lut

    @timing
    def _lazy_load_luts(self, luts_names):
        """
        lazy load luts from xml files
        Parameters
        ----------
        luts_names: list of str

        Returns
        -------
        xarray.Dataset with variables from `luts_names`.

        """

        luts_list = []
        luts = None
        for lut_name in luts_names:
            xml_type = self._map_lut_files[lut_name]
            xml_files = self.s1meta.files.copy()
            # polarization is a category. we use codes (ie index),
            # to have well ordered polarizations in latter combine_by_coords
            xml_files['pol_code'] = xml_files['polarization'].cat.codes
            xml_files = xml_files.set_index('pol_code')[xml_type]

            if not self._vars_with_pol[lut_name]:
                # luts are identical in all pols: take the fist one
                xml_files = xml_files.iloc[[0]]

            for pol_code, xml_file in xml_files.iteritems():
                pol = self.s1meta.files['polarization'].cat.categories[pol_code]
                if self._vars_with_pol[lut_name]:
                    name = "%s_%s" % (lut_name, pol)
                else:
                    name = lut_name

                # get the lut function. As it takes some time to parse xml, make it delayed
                lut_f_delayed = dask.delayed(self.s1meta.xml_parser.get_compound_var)(xml_file, lut_name)
                lut = map_blocks_coords(
                    self._da_tmpl.astype(self._dtypes[lut_name]),
                    lut_f_delayed,
                    name='blocks_%s' % name
                )

                # needs to add pol dim ?
                if self._vars_with_pol[lut_name]:
                    lut = lut.assign_coords(pol_code=pol_code).expand_dims('pol_code')

                # set xml file and xpath used as history
                histo = self.s1meta.xml_parser.get_compound_var(xml_file, lut_name,
                                                                describe=True)
                lut.name = lut_name
                if self._patch_variable:
                    lut = self._patch_lut(lut)
                lut.attrs['history'] = histo
                lut = lut.to_dataset()

                luts_list.append(lut)
            luts = xr.combine_by_coords(luts_list)

            # convert pol_code to string
            pols = self.s1meta.files['polarization'].cat.categories[luts.pol_code.values.tolist()]
            luts = luts.rename({'pol_code': 'pol'}).assign_coords({'pol': pols})
        return luts

    @timing
    def _load_digital_number(self, resolution=None, chunks=None, resampling=rasterio.enums.Resampling.rms):
        """
        load digital_number from self.s1meta.files['measurement'], as an `xarray.Dataset`.

        Parameters
        ----------
        resolution: None, number, str or dict
            see `xsar.open_dataset`
        resampling: rasterio.enums.Resampling
            see `xsar.open_dataset`

        Returns
        -------
        xarray.Dataset
            dataset (possibly dual-pol), with basic coords/dims naming convention
        """

        map_dims = {
            'pol': 'band',
            'atrack': 'y',
            'xtrack': 'x'
        }

        if resolution is not None:
            comment = 'resampled at "%s" with %s.%s.%s' % (
                resolution, resampling.__module__, resampling.__class__.__name__, resampling.name)
        else:
            comment = 'read at full resolution'

        # arbitrary rio object, to get shape, etc ... (will not be used to read data)
        rio = rasterio.open(self.s1meta.files['measurement'].iloc[0])

        chunks['pol'] = 1
        # sort chunks keys like map_dims
        chunks = dict(sorted(chunks.items(), key=lambda pair: list(map_dims.keys()).index(pair[0])))
        chunks_rio = {map_dims[d]: chunks[d] for d in map_dims.keys()}

        if resolution is None:
            # using tiff driver: need to read individual tiff and concat them
            # riofiles['rio'] is ordered like self.s1meta.manifest_attrs['polarizations']

            dn = xr.concat(
                [
                    rioxarray.open_rasterio(
                        f, chunks=chunks_rio, parse_coordinates=False
                    ) for f in self.s1meta.files['measurement']
                ], 'band'
            ).assign_coords(band=np.arange(len(self.s1meta.manifest_attrs['polarizations'])) + 1)

            # set dimensions names
            dn = dn.rename(dict(zip(map_dims.values(), map_dims.keys())))

            # create coordinates from dimension index (because of parse_coordinates=False)
            dn = dn.assign_coords({'atrack': dn.atrack, 'xtrack': dn.xtrack})
            dn = dn.drop_vars('spatial_ref', errors='ignore')
        else:
            if not isinstance(resolution, dict):
                if isinstance(resolution, str) and resolution.endswith('m'):
                    resolution = float(resolution[:-1])
                resolution = dict(atrack=resolution / self.s1meta.pixel_atrack_m,
                                  xtrack=resolution / self.s1meta.pixel_xtrack_m)

            # resample the DN at gdal level, before feeding it to the dataset
            out_shape = (
                int(rio.height / resolution['atrack']),
                int(rio.width / resolution['xtrack'])
            )
            out_shape_pol = (1,) + out_shape
            # read resampled array in one chunk, and rechunk
            # this doesn't optimize memory, but total size remain quite small

            if isinstance(resolution['atrack'], int):
                # legacy behaviour: winsize is the maximum full image size that can be divided  by resolution (int)
                winsize = (0, 0, rio.width // resolution['xtrack'] * resolution['xtrack'],
                           rio.height // resolution['atrack'] * resolution['atrack'])
                window = rasterio.windows.Window(*winsize)
            else:
                window = None

            dn = xr.concat(
                [
                    xr.DataArray(
                        dask.array.from_array(
                            rasterio.open(f).read(
                                out_shape=out_shape_pol,
                                resampling=resampling,
                                window=window
                            ),
                            chunks=chunks_rio
                        ),
                        dims=tuple(map_dims.keys()), coords={'pol': [pol]}
                    ) for f, pol in
                    zip(self.s1meta.files['measurement'], self.s1meta.manifest_attrs['polarizations'])
                ],
                'pol'
            ).chunk(chunks)

            # create coordinates at box center
            translate = Affine.translation((resolution['xtrack'] - 1) / 2, (resolution['atrack'] - 1) / 2)
            scale = Affine.scale(
                rio.width // resolution['xtrack'] * resolution['xtrack'] / out_shape[1],
                rio.height // resolution['atrack'] * resolution['atrack'] / out_shape[0])
            xtrack, _ = translate * scale * (dn.xtrack, 0)
            _, atrack = translate * scale * (0, dn.atrack)
            dn = dn.assign_coords({'atrack': atrack, 'xtrack': xtrack})

        # for GTiff driver, pols are already ordered. just rename them
        dn = dn.assign_coords(pol=self.s1meta.manifest_attrs['polarizations'])

        if not all(self.s1meta.denoised.values()):
            descr = 'denoised'
        else:
            descr = 'not denoised'
        var_name = 'digital_number'

        dn.attrs = {
            'comment': '%s digital number, %s' % (descr, comment),
            'history': yaml.safe_dump(
                {
                    var_name: get_glob(
                        [p.replace(self.s1meta.path + '/', '') for p in self.s1meta.files['measurement']])
                }
            )
        }
        ds = dn.to_dataset(name=var_name)
        astype = self._dtypes.get(var_name)
        if astype is not None:
            ds = ds.astype(self._dtypes[var_name])

        return ds

    @timing
    def _load_from_geoloc(self, varnames):
        """
        Interpolate (with RectBiVariateSpline) variables from `self.s1meta.geoloc` to `self._dataset`

        Parameters
        ----------
        varnames: list of str
            subset of variables names in `self.s1meta.geoloc`

        Returns
        -------
        xarray.Dataset
            With interpolated vaiables

        """

        da_list = []

        def interp_func_slc(vect1dazti, vect1dxtrac, **kwargs):
            """

            Parameters
            ----------
            vect1dazti (np.ndarray) : azimuth times at high resolution
            vect1dxtrac (np.ndarray): range coords

            Returns
            -------

            """
            # exterieur de boucle
            rbs = kwargs['rbs']

            def wrapperfunc(*args, **kwargs):
                rbs2 = args[2]
                return rbs2(args[0], args[1], grid=False)

            return wrapperfunc(vect1dazti[:, np.newaxis], vect1dxtrac[np.newaxis, :], rbs)

        for varname in varnames:
            logger.debug('varname : %s', varname)
            if varname in ['azimuth_time']:
                z_values = self.s1meta.geoloc[varname].astype(float)
            elif varname == 'longitude':
                z_values = self.s1meta.geoloc[varname]
                if self.s1meta.cross_antemeridian:
                    logger.debug('translate longitudes between 0 and 360')
                    z_values = z_values % 360
            else:
                z_values = self.s1meta.geoloc[varname]
            if self.s1meta._bursts['burst'].size != 0:
                # TOPS SLC
                rbs = RectBivariateSpline(
                    self.s1meta.geoloc.azimuth_time[:, 0].astype(float),
                    self.s1meta.geoloc.xtrack,
                    z_values,
                    kx=1, ky=1,
                )
                interp_func = interp_func_slc
            else:
                rbs = None
                interp_func = RectBivariateSpline(
                    self.s1meta.geoloc.atrack,
                    self.s1meta.geoloc.xtrack,
                    z_values,
                    kx=1, ky=1
                )
            # the following take much cpu and memory, so we want to use dask
            # interp_func(self._dataset.atrack, self.dataset.xtrack)
            typee = self.s1meta.geoloc[varname].dtype
            if self.s1meta._bursts['burst'].size != 0:
                datemplate = self._da_tmpl.astype(typee).copy()
                # replace the atrack coordinates by atrack_time coordinates
                datemplate = datemplate.assign_coords({'atrack': datemplate.coords['atrack_time']})
                da_var = map_blocks_coords(
                    datemplate,
                    interp_func,
                    func_kwargs={"rbs": rbs}
                )
                # put back the real atrack coordinates
                da_var = da_var.assign_coords({'atrack': self._dataset.digital_number.atrack})
            else:
                da_var = map_blocks_coords(
                    self._da_tmpl.astype(typee),
                    interp_func
                )
            if varname == 'longitude':
                if self.s1meta.cross_antemeridian:
                    da_var.data = da_var.data.map_blocks(to_lon180)

            da_var.name = varname

            # copy history
            try:
                da_var.attrs['history'] = self.s1meta.geoloc[varname].attrs['history']
            except KeyError:
                pass

            da_list.append(da_var)

        return xr.merge(da_list)

    @timing
    def _load_ground_heading(self):
        def coords2heading(atracks, xtracks):
            return self.s1meta.coords2heading(atracks, xtracks, to_grid=True, approx=True)

        gh = map_blocks_coords(
            self._da_tmpl.astype(self._dtypes['ground_heading']),
            coords2heading,
            name='ground_heading'
        )

        gh.attrs = {
            'comment': 'at ground level, computed from lon/lat in atrack direction'
        }

        return gh.to_dataset(name='ground_heading')

    @timing
    def _load_rasterized_masks(self):
        def _test(atrack, xtrack, mask=None):
            chunk_coords = bbox_coords(atrack, xtrack, pad=None)
            # chunk footprint polygon, in dataset coordinates (with buffer, to enlarge a little the footprint)
            chunk_footprint_coords = Polygon(chunk_coords).buffer(10)
            tmp = self.s1meta.name
            # vector_mask_ll = s1meta.get_mask(mask) #
            return np.meshgrid(xtrack, atrack)[0] * 0

        def _rasterize_mask_by_chunks(atrack, xtrack, mask='land'):
            chunk_coords = bbox_coords(atrack, xtrack, pad=None)
            # chunk footprint polygon, in dataset coordinates (with buffer, to enlarge a little the footprint)
            chunk_footprint_coords = Polygon(chunk_coords).buffer(10)
            # chunk footprint polygon, in lon/lat
            chunk_footprint_ll = self.s1meta.coords2ll(chunk_footprint_coords)

            # get vector mask over chunk, in lon/lat
            vector_mask_ll = self.s1meta.get_mask(mask).intersection(chunk_footprint_ll)

            if vector_mask_ll.is_empty:
                # no intersection with mask, return zeros
                return np.zeros((atrack.size, xtrack.size))

            # vector mask, in atrack/xtrack coordinates
            vector_mask_coords = self.s1meta.ll2coords(vector_mask_ll)

            # shape of the returned chunk
            out_shape = (atrack.size, xtrack.size)

            # transform * (x, y) -> (atrack, xtrack)
            # (where (x, y) are index in out_shape)
            # Affine.permutation() is used because (atrack, xtrack) is transposed from geographic

            transform = Affine.translation(*chunk_coords[0]) * Affine.scale(
                *[np.unique(np.diff(c))[0] for c in [atrack, xtrack]]) * Affine.permutation()

            raster_mask = rasterio.features.rasterize(
                [vector_mask_coords],
                out_shape=out_shape,
                all_touched=False,
                transform=transform
            )
            return raster_mask

        da_list = []
        for mask in self.s1meta.mask_names:
            da_mask = map_blocks_coords(
                self._da_tmpl,
                _rasterize_mask_by_chunks,
                func_kwargs={'mask': mask}
            )
            name = '%s_mask' % mask
            da_mask.attrs['history'] = yaml.safe_dump({name: self.s1meta.get_mask(mask, describe=True)})
            da_list.append(da_mask.to_dataset(name=name))

        return xr.merge(da_list)

    @timing
    def _load_rasters_vars(self):
        # load and map variables from rasterfile (like ecmwf) on dataset
        if self.s1meta.rasters.empty:
            return None
        else:
            logger.warning('Raster variable are experimental')

        if self.s1meta.cross_antemeridian:
            raise NotImplementedError('Antimeridian crossing not yet checked')

        # get lon/lat box for xsar dataset
        lons, lats = list(zip(*self.s1meta.footprint.exterior.coords))
        lon_range = [min(lons), max(lons)]
        lat_range = [min(lats), max(lats)]

        # will contain xr.DataArray to merge
        da_var_list = []

        for name, infos in self.s1meta.rasters.iterrows():
            # read the raster file using helpers functions
            read_function = infos['read_function']
            get_function = infos['get_function']
            resource = infos['resource']

            kwargs = {
                's1meta': self,
                'date': self.s1meta.start_date,
                'footprint': self.s1meta.footprint
            }

            logger.debug('adding raster "%s" from resource "%s"' % (name, str(resource)))
            if get_function is not None:
                try:
                    resource_dec = get_function(resource, **kwargs)
                except TypeError:
                    resource_dec = get_function(resource)

            if read_function is None:
                raster_ds = xr.open_dataset(resource_dec, chunk=1000)
            else:
                # read_function should return a chunked dataset (so it's fast)
                raster_ds = read_function(resource_dec)

            # add globals raster attrs to globals dataset attrs
            hist_res = {'resource': resource}
            if get_function is not None:
                hist_res.update({'resource_decoded': resource_dec})

            if not raster_ds.rio.crs.is_geographic:
                raise NotImplementedError("Non geographic crs not implemented")

            # ensure dim ordering
            raster_ds = raster_ds.transpose('y', 'x')
            if np.all(raster_ds.y.diff('y') <= 0):
                # sort y (lat) ascending (for RectBiVariateSpline)
                raster_ds = raster_ds.reindex(y=raster_ds.y[::-1])

            # from lon/lat box in xsar dataset, get the corresponding box in raster_ds (by index)
            ilon_range = [
                np.searchsorted(raster_ds.x.values, lon_range[0], side='right'),
                np.searchsorted(raster_ds.x.values, lon_range[1], side='left')
            ]
            ilat_range = [
                np.searchsorted(raster_ds.y.values, lat_range[0], side='right'),
                np.searchsorted(raster_ds.y.values, lat_range[1], side='left')
            ]
            # select the xsar box in the raster
            raster_ds = raster_ds.isel(x=slice(*ilon_range), y=slice(*ilat_range))

            # 1D array of lons/lats, trying to have same spacing as dataset (if not to high)
            num = min((self._dataset.xtrack.size + self._dataset.atrack.size) // 2, 1000)
            lons = np.linspace(*lon_range, num=num)
            lats = np.linspace(*lat_range, num=num)

            @dask.delayed
            def _map_raster2xsar(da):
                # map the 'da' dataarray variable from the raster to xsar dataset
                da = da.drop_vars(['spatial_ref', 'crs'], errors='ignore')

                upscaled_da = map_blocks_coords(
                    xr.DataArray(dims=['y', 'x'], coords={'x': lons, 'y': lats}).chunk(1000),
                    RectBivariateSpline(da.y.values, da.x.values, da.values)
                )

                reprojected_da = upscaled_da.interp(
                    x=self._dataset.longitude,
                    y=self._dataset.latitude
                )
                reprojected_da = reprojected_da.drop_vars(['x', 'y', 'spatial_ref', 'crs'], errors='ignore')
                reprojected_da.attrs.update(da.attrs)

                reprojected_da.name = '%s_%s' % (name, da.name)

                # reprojected_da has same shape as other variables is xsar dataset, with optional 3rd dim
                return reprojected_da

            for var in raster_ds:
                var_name = '%s_%s' % (name, raster_ds[var].name)
                da_var = xr.DataArray(
                    dask.array.from_delayed(
                        _map_raster2xsar(raster_ds[var]),
                        self._da_tmpl.shape,
                        dtype='f8',
                        name='%s' % var_name
                    ),
                    dims=['atrack', 'xtrack'],
                    coords={'atrack': self._da_tmpl.atrack, 'xtrack': self._da_tmpl.xtrack},
                    attrs=raster_ds[var].attrs
                ).chunk(self._da_tmpl.chunks)
                da_var.attrs['history'] = yaml.safe_dump({var_name: hist_res})
                logger.debug('adding variable "%s" from raster "%s"' % (var_name, name))
                da_var_list.append(da_var)

        return xr.merge(da_var_list)

    def _get_lut(self, var_name):
        """
        Get lut for `var_name`

        Parameters
        ----------
        var_name: str

        Returns
        -------
        xarray.Dataarray
            lut for `var_name`
        """
        try:
            lut_name = self._map_var_lut[var_name]
        except KeyError:
            raise ValueError("can't find lut name for var '%s'" % var_name)
        try:
            lut = self._luts[lut_name]
        except KeyError:
            raise ValueError("can't find lut from name '%s' for variable '%s' " % (lut_name, var_name))
        return lut

    def _apply_calibration_lut(self, var_name):
        """
        Apply calibration lut to `digital_number` to compute `var_name`.
        see https://sentinel.esa.int/web/sentinel/radiometric-calibration-of-level-1-products

        Parameters
        ----------
        var_name: str
            Variable name to compute by applying lut. Must exist in `self._map_var_lut`` to be able to get the corresponding lut.

        Returns
        -------
        xarray.Dataset
            with one variable named by `var_name`
        """
        lut = self._get_lut(var_name)
        res = (np.abs(self._dataset.digital_number) ** 2. / (lut ** 2))
        # dn default value is 0: convert to Nan
        res = res.where(res > 0)
        astype = self._dtypes.get(var_name)
        if astype is not None:
            res = res.astype(astype)

        res.attrs.update(lut.attrs)
        res.attrs['history'] = merge_yaml([lut.attrs['history']], section=var_name)
        res.attrs['references'] = 'https://sentinel.esa.int/web/sentinel/radiometric-calibration-of-level-1-products'

        return res.to_dataset(name=var_name)

    def reverse_calibration_lut(self, ds_var):
        """
        TODO: replace ds_var by var_name
        Inverse of `_apply_calibration_lut` : from `var_name`, reverse apply lut, to get digital_number.
        See `official ESA documentation <https://sentinel.esa.int/web/sentinel/radiometric-calibration-of-level-1-products>`_ .
        > Level-1 products provide four calibration Look Up Tables (LUTs) to produce ß0i, σ0i and γi
        > or to return to the Digital Number (DN)

        A warning message may be issued if original complex 'digital_number' is converted to module during this operation.

        Parameters
        ----------
        ds_var: xarray.Dataset
            with only one variable name that must exist in `self._map_var_lut` to be able to reverse the lut to get digital_number

        Returns
        -------
        xarray.Dataset
            with one variable named 'digital_number'.
        """
        var_names = list(ds_var.keys())
        assert len(var_names) == 1
        var_name = var_names[0]
        if var_name not in self._map_var_lut:
            raise ValueError(
                "Unable to find lut for var '%s'. Allowed : %s" % (var_name, str(self._map_var_lut.keys())))
        da_var = ds_var[var_name]
        lut = self._luts[self._map_var_lut[var_name]]

        # resize lut with same a/xtrack as da_var
        lut = lut.sel(atrack=da_var.atrack, xtrack=da_var.xtrack, method='nearest')
        # as we used 'nearest', force exact coords
        lut['atrack'] = da_var.atrack
        lut['xtrack'] = da_var.xtrack
        # waiting for https://github.com/pydata/xarray/pull/4155
        # lut = lut.interp(atrack=da_var.atrack, xtrack=da_var.xtrack)

        # revert lut to get dn
        dn = np.sqrt(da_var * lut ** 2)

        if self._dataset.digital_number.dtype == np.complex and dn.dtype != np.complex:
            warnings.warn(
                "Unable to retrieve 'digital_number' as dtype '%s'. Fallback to '%s'"
                % (str(self._dataset.digital_number.dtype), str(dn.dtype))
            )

        name = 'digital_number'
        ds = dn.to_dataset(name=name)

        return ds

    def _get_noise(self, var_name):
        """
        Get noise equivalent for  `var_name`.
        see https://sentinel.esa.int/web/sentinel/radiometric-calibration-of-level-1-products

        Parameters
        ----------
        var_name: str
            Variable name to compute. Must exist in `self._map_var_lut` to be able to get the corresponding lut.

        Returns
        -------
        xarray.Dataset
            with one variable named by `'ne%sz' % var_name[0]` (ie 'nesz' for 'sigma0', 'nebz' for 'beta0', etc...)
        """
        noise_lut = self._luts['noise_lut']
        lut = self._get_lut(var_name)
        dataarr = noise_lut / lut ** 2
        name = 'ne%sz' % var_name[0]
        astype = self._dtypes.get(name)
        if astype is not None:
            dataarr = dataarr.astype(astype)
        dataarr.attrs['history'] = merge_yaml([lut.attrs['history'], noise_lut.attrs['history']], section=name)
        return dataarr.to_dataset(name=name)

    def _add_denoised(self, ds, clip=False, vars=None):
        """add denoised vars to dataset

        Parameters
        ----------
        ds : xarray.DataSet
            dataset with non denoised vars, named `%s_raw`.
        clip : bool, optional
            If True, negative signal will be clipped to 0. (default to False )
        vars : list, optional
            variables names to add, by default `['sigma0' , 'beta0' , 'gamma0']`

        Returns
        -------
        xarray.DataSet
            dataset with denoised vars
        """
        if vars is None:
            vars = ['sigma0', 'beta0', 'gamma0']
        for varname in vars:
            varname_raw = varname + '_raw'
            noise = 'ne%sz' % varname[0]
            if varname_raw not in ds:
                continue
            if all(self.s1meta.denoised.values()):
                # already denoised, just add an alias
                ds[varname] = ds[varname_raw]
            elif len(set(self.s1meta.denoised.values())) != 1:
                # TODO: to be implemented
                raise NotImplementedError("semi denoised products not yet implemented")
            else:
                denoised = ds[varname_raw] - ds[noise]
                denoised.attrs['history'] = merge_yaml(
                    [ds[varname_raw].attrs['history'], ds[noise].attrs['history']],
                    section=varname
                )
                if clip:
                    denoised = denoised.clip(min=0)
                    denoised.attrs['comment'] = 'clipped, no values <0'
                else:
                    denoised.attrs['comment'] = 'not clipped, some values can be <0'
                ds[varname] = denoised
        return ds

    @property
    def _burst_azitime(self):
        """
        Get azimuth time at high resolution.

        Returns
        -------
        xarray.DataArray
            the high resolution azimuth time vector interpolated at the middle of the sub-swath
        """
        azitime = self.s1meta._burst_azitime()
        iz = np.searchsorted(azitime.atrack, self._dataset.atrack)
        azitime = azitime.isel({'atrack': iz})
        azitime = azitime.assign_coords({"atrack": self._dataset.atrack})
        return azitime

    def _get_sensor_velocity(self):
        """
        Interpolated sensor velocity
        Returns
        -------
        xarray.Dataset()
            containing a single variable velocity
        """

        azimuth_times = self._burst_azitime
        orbstatevect = self.s1meta.orbit
        azi_times = orbstatevect.index.values
        velos = np.array([[uu.x ** 2., uu.y ** 2., uu.z ** 2.] for uu in orbstatevect['velocity'].values])
        vels = np.sqrt(np.sum(velos, axis=1))
        interp_f = interp1d(azi_times.astype(float), vels)
        _vels = interp_f(azimuth_times.astype(float))
        res = xr.DataArray(_vels, dims=['atrack'], coords={'atrack': self.dataset.atrack})
        return xr.Dataset({'velocity': res})

    def _range_ground_spacing(self):
        """
        Get SAR image range ground spacing.

        Parameters
        ----------
        Returns
        -------
        range_ground_spacing_vect : xarray.DataArray
            range ground spacing (xtrack coordinates)

        Notes
        -----
        For GRD products is it the same same value along xtrack axis
        """
        ground_spacing = np.array(self.s1meta.image['slant_pixel_spacing'])
        if self.s1meta.product == 'SLC':
            atrack_tmp = self._dataset['atrack']
            xtrack_tmp = self._dataset['xtrack']
            # get the incidence at the middle of atrack dimension of the part of image selected
            inc = self._dataset['incidence'].isel({'atrack': int(len(atrack_tmp) / 2),
                                                   })
            range_ground_spacing_vect = ground_spacing[1] / np.sin(np.radians(inc))
            range_ground_spacing_vect.attrs['history'] = ''

        else:  # GRD
            valuess = np.ones((len(self._dataset['xtrack']))) * ground_spacing[1]
            range_ground_spacing_vect = xr.DataArray(valuess, coords={'xtrack': self._dataset['xtrack']},
                                                     dims=['xtrack'])
        return xr.Dataset({'range_ground_spacing': range_ground_spacing_vect})

    def __repr__(self):
        if self.sliced:
            intro = "sliced"
        else:
            intro = "full covevage"
        return "<Sentinel1Dataset %s object>" % intro

    def _repr_mimebundle_(self, include=None, exclude=None):
        return repr_mimebundle(self, include=include, exclude=exclude)

import numpy,from numpy,from scipy,from cupy


#/usr/bin/python

'''

This function collects NEXMD populations from surface hopping
and populations from Schrodinger equation.

Two files are generated in the current working directory if
collecting NEXMD populations are requested, 'pop.out' and
'pop.err'.  In 'pop.out', first column is time in fs, followed by
the average population on each PES, followed by the sum of all
PES populations (should be = 1.0), followed by the average
populations from quantum coefficients, followed by the sum of
quantum populations (should be = 1.0).  In 'pop.err', first
column is directory of trajectory, third column is the time at
which the trajectory has ended in fs, or 'does not exist'.  The
'pop.err' file is not generated if all trajectories are complete.

'Completed' trajectories are trajectories that have completed
within the time defined by user while executing this function.
'Excellent' trajectories are trajectories that have completed
within the time defined in 'header' located in the NEXMD
directory.

Output Files:
- pop_[type].out, where [type] = mean_ensemble

Error Files:
- pop_[type].err, where [type] = mean_ensemble

'''

import numpy as np
import os
import sys
import glob

cwd = os.getcwd()

def population(header):

    print 'Collecting populations.'

    ## Directory names ##
    NEXMDir = raw_input('NEXMD directory: ')
    if not os.path.exists(NEXMDir):
        print 'Path %s does not exist.' % (NEXMDir)
        sys.exit()
    ## Check if NEXMD folders exist ##
    NEXMDs = glob.glob('%s/NEXMD*/' % (NEXMDir))
    NEXMDs.sort()
    if len(NEXMDs) == 0:
        print 'There are no NEXMD folders in %s.' % (NEXMDir)
        sys.exit()

    ## Information from header ##
    if not os.path.exists('%s/header' % (NEXMDir)):
        print 'Path %s/header does not exist.' % (NEXMDir)
        sys.exit()
    header = header('%s/header' % (NEXMDir))

    ## Adding + 1 to include zeroth time-step ##
    header.n_class_steps = header.n_class_steps + 1

    ## Collection time ##
    tcoll = input('Calculate populations up to what time in femtoseconds?\nNote that averaged results will only include trajectories that are complete up to this time: ')
    if isinstance(tcoll, int) == False and isinstance(tcoll, float) == False:
        print 'Time must be integer or float.'
        sys.exit()
    if tcoll < 0:
        print 'Time must be integer or float greater than zero.'
        sys.exit()
    tcoll = np.float(tcoll)
    if tcoll > (header.n_class_steps - 1)*header.time_step:
        tcoll = (header.n_class_steps -1)*header.time_step

    ## Number of classical time-steps ##
    tscol = 0
    while tscol*header.time_step*header.out_data_steps <= tcoll:
        tscol += 1

    ## Collection time array ##
    times = np.around(np.linspace(header.time_init, tcoll, tscol), decimals = 3)
    fpoph = np.zeros((tscol,header.n_exc_states_propagate))
    fpopc = np.zeros((tscol,header.n_exc_states_propagate))

    ## Collect populations ##
    output = open('%s/pop_mean_ensemble.out' % (cwd),'w')
    error = open('%s/pop_mean_ensemble.err' % (cwd),'w')
    ttraj = 0
    ctraj = 0
    etraj = 0
    errflag = 0
    for NEXMD in NEXMDs:
        if not os.path.exists('%s/dirlist1' % (NEXMD)):
            print 'Path %sdirlist1 does not exist.' % (NEXMD)
            sys.exit()
        dirlist1 = np.int_(np.genfromtxt('%s/dirlist1' % (NEXMD)))
        if isinstance(dirlist1,int) == True:
            dirlist1 = np.array([dirlist1])
        for dir in dirlist1:
            if not os.path.exists('%s/%04d/coeff-n.out' % (NEXMD,dir)):
                print >> error, 'Path %s%04d/coeff-n.out does not exist.' % (NEXMD,dir)
                errflag = 1
                ttraj += 1
                continue
            data = open('%s/%04d/coeff-n.out' % (NEXMD,dir),'r')
            data = data.readlines()
            tsteps = len(data)
            tflag = 0
            if tsteps >= tscol:
                poph = np.zeros((tscol,header.n_exc_states_propagate))
                popc = np.zeros((tscol,header.n_exc_states_propagate))
                index = 0
                for line in data[0:tscol:1]:
                    val = line.split()
                    pes = np.int(val[0])
                    time = np.around(np.float(val[1]), decimals = 3)
                    if time != times[index]:
                        print >> error, 'There is an inconsistency in time-step in %s%04d at %.3f fs' % (NEXMD,dir,times[index])
                        tflag = 1
                        errflag = 1
                        break
                    poph[index][pes-1] = 1.0
                    popc[index] = np.float_(val[2:2+header.n_exc_states_propagate])
                    index += 1
                if tflag == 0:
                    fpoph += poph
                    fpopc += popc
                    print '%s%04d' % (NEXMD,dir), '%0*.2f' % (len(str((header.n_class_steps))) + 2, (tsteps - 1)*header.time_step)
                    ctraj += 1
                    if tsteps == header.n_class_steps:
                        etraj += 1
            else:
                print '%s%04d' % (NEXMD,dir), '%0*.2f' % (len(str((header.n_class_steps))) + 2, (tsteps - 1)*header.time_step)
                print >> error, '%s%04d' % (NEXMD,dir), '%0*.2f' % (len(str((header.n_class_steps))) + 2, (tsteps - 1)*header.time_step)
                errflag = 1
            ttraj += 1
    if ctraj == 0:
        print 'No trajectories completed within %0*.2f fs.' % (len(str(header.n_class_steps)),tcoll)
        os.remove('%s/pop_mean_ensemble.out' % (cwd))
    else:
        fpoph = fpoph/ctraj
        fpopc = fpopc/ctraj
        print 'Total trajectories:', '%04d' % (ttraj)
        print 'Completed trajectories:', '%04d' % (ctraj)
        print 'Excellent trajectories:', '%04d' % (etraj)
        print >> output, 'Total trajectories: ', '%04d' % (ttraj)
        print >> output, 'Completed trajectories: ', '%04d' % (ctraj)
        print >> output, 'Excellent trajectories: ', '%04d' % (etraj)
        for tstep in np.arange(tscol):
            print >> output, '%0*.2f' % (len(str((header.n_class_steps))) + 2,header.time_step*tstep), ' '.join(str('%.3f' % (x)) for x in fpoph[tstep]), '%.3f' % (np.sum(fpoph[tstep])), ' '.join(str('%.3f' % (x)) for x in fpopc[tstep]), '%.3f' % (np.sum(fpopc[tstep]))
    if errflag == 1:
        print 'One or more trajectories have experienced an error, check pop_mean_ensemble.err.'
    else:
        os.remove('%s/pop_mean_ensemble.err' % (cwd))