from .ransac import ransac
from .utils import dist_matrix, orientation_diff
import numpy as np
import torch

def select_seeds(dist1: torch.Tensor, R1: float, scores1: torch.Tensor, fnn12: torch.Tensor, mnn: torch.Tensor):
    """
        Select seed correspondences among the set of available matches.

        dist1: Precomputed distance matrix between keypoints in image I_1
        R1: Base radius of neighborhoods in image I_1
        scores1: Confidence scores on the putative_matches. Usually holds Lowe's ratio scores.
        fnn12: Matches between keypoints of I_1 and I_2.
               The i-th entry of fnn12 is j if and only if keypoint k_i in image I_1 is matched to keypoint k_j in image I_2
        mnn: A mask indicating which putative matches are also mutual nearest neighbors. See documentation on 'force_seed_mnn' in the DEFAULT_CONFIG.
             If None, it disables the mutual nearest neighbor filtering on seed point selection.
             Expected a bool tensor with shape (num_keypoints_in_source_image,)

        Returns:
            Indices of seed points.

            im1seeds: Keypoint index of chosen seeds in image I_1
            im2seeds: Keypoint index of chosen seeds in image I_2
    """
    im1neighmap = dist1 < R1**2  # (n1, n1)
    # find out who scores higher than whom
    im1scorescomp = scores1.unsqueeze(1) > scores1.unsqueeze(0)  # (n1, n1)
    # find out who scores higher than all of its neighbors: seed points
    if mnn is not None:
        im1bs = (~torch.any(im1neighmap & im1scorescomp & mnn.unsqueeze(0),
                            dim=1)) & mnn & (scores1 < 0.8**2)  # (n1,)
    else:
        im1bs = (~torch.any(im1neighmap & im1scorescomp, dim=1)) & (scores1 <
                                                                    0.8**2)

    # collect all seeds in both images and the 1NN of the seeds of the other image
    im1seeds = torch.where(im1bs)[0]  # (n1bs) index format
    im2seeds = fnn12[im1bs]  # (n1bs) index format
    return im1seeds, im2seeds

def extract_neighborhood_sets(
        o1: torch.Tensor, o2: torch.Tensor, s1: torch.Tensor, s2: torch.Tensor,
        dist1: torch.Tensor, im1seeds: torch.Tensor, im2seeds: torch.Tensor,
        k1: torch.Tensor, k2: torch.Tensor, R1: float, R2: float,
        fnn12: torch.Tensor, ORIENTATION_THR: float, SCALE_RATE_THR: float,
        SEARCH_EXP: float, MIN_INLIERS: float):
    """
        Assign keypoints to seed points. This checks both the distance and
        the agreement of the local transformation if available.

        o1: Orientations of keypoints in image I_1
        o2: Orientations of keypoints in image I_2
        s1: Scales of keypoints in image I_1
        s2: Scales of keypoints in image I_2
        dist1: Precomputed distance matrix between keypoints in image I_1
        im1seeds: Keypoint index of chosen seeds in image I_1
        im2seeds: Keypoint index of chosen seeds in image I_2
        k1: Keypoint locations in image I_1
        k2: Keypoint locations in image I_2
        R1: Base radius of neighborhoods in image I_1
        R2: Base radius of neighborhoods in image I_2
        fnn12: Matches between keypoints of I_1 and I_2.
               The i-th entry of fnn12 is j if and only if keypoint k_i in image I_1 is matched to keypoint k_j in image I_2
        ORIENTATION_THR: Maximum deviation of orientation with respect to seed S_i to keep a keypoint in i-th neighborhood
        SCALE_RATE_THR: Maximum deviation of scale with respect to seed S_i to keep a keypoint in i-th neighborhood
        SEARCH_EXP: Expansion rate for both radii R1 and R2 to consider inclusion of neighboring keypoints
        MIN_INLIERS: Minimum number of inliers to keep a seed point. This is used as an early filter here
                     to remove already seeds with not enough samples to ever pass this threshold.

        Returns:
            Local neighborhoods assignments:

            local_neighs_mask: Boolean matrix of size (num_seeds, num_keypoints).
                               Entry (i, j) is True iff keypoint j was assigned to seed i.
            rdims: Number of keypoints included in the neighborhood for each seed
            im1seeds: Keypoint index of chosen seeds in image I_1
            im2seeds: Keypoint index of chosen seeds in image I_2

    """
    dst1 = dist1[im1seeds, :]
    dst2 = dist_matrix(k2[fnn12[im1seeds]], k2[fnn12])

    # initial candidates are matches which are close to the same seed in both images
    local_neighs_mask = (dst1 < (SEARCH_EXP * R1) ** 2) \
                        & (dst2 < (SEARCH_EXP * R2) ** 2)

    # If requested, also their orientation delta should be compatible with that of the corresponding seed
    if ORIENTATION_THR is not None and ORIENTATION_THR < 180:
        relo = orientation_diff(o1, o2[fnn12])
        orientation_diffs = torch.abs(
            orientation_diff(relo.unsqueeze(0), relo[im1seeds].unsqueeze(1)))
        local_neighs_mask = local_neighs_mask & (orientation_diffs <
                                                 ORIENTATION_THR)

    # If requested, also their scale delta should be compatible with that of the corresponding seed
    if SCALE_RATE_THR is not None and SCALE_RATE_THR < 10:
        rels = s2[fnn12] / s1
        scale_rates = rels[im1seeds].unsqueeze(1) / rels.unsqueeze(0)
        local_neighs_mask = local_neighs_mask & (scale_rates < SCALE_RATE_THR) \
                            & (scale_rates > 1 / SCALE_RATE_THR)  # (ns, n1)

    # count how many keypoints ended up in each neighborhood
    numn1 = torch.sum(local_neighs_mask, dim=1)
    # and only keep the ones that have enough points
    valid_seeds = numn1 >= MIN_INLIERS

    local_neighs_mask = local_neighs_mask[valid_seeds, :]

    rdims = numn1[valid_seeds]

    return local_neighs_mask, rdims, im1seeds[valid_seeds], im2seeds[
        valid_seeds]

def extract_local_patterns(
        fnn12: torch.Tensor,
        fnn_to_seed_local_consistency_map_corr: torch.Tensor, k1: torch.Tensor,
        k2: torch.Tensor, im1seeds: torch.Tensor, im2seeds: torch.Tensor,
        scores: torch.Tensor):
    """
        Prepare local neighborhoods around each seed for the parallel RANSACs. This involves two steps:
            1) Collect all selected keypoints and refer them with respect to their seed point
            2) Sort keypoints by score for the progressive sampling to pick the best samples first

        fnn12: Matches between keypoints of I_1 and I_2.
               The i-th entry of fnn12 is j if and only if keypoint k_i in image I_1 is matched to keypoint k_j in image I_2
        fnn_to_seed_local_consistency_map_corr: Boolean matrix of size (num_seeds, num_keypoints).
                                                Entry (i, j) is True iff keypoint j was assigned to seed i.
        k1: Keypoint locations in image I_1
        k2: Keypoint locations in image I_2
        im1seeds: Keypoint index of chosen seeds in image I_1
        im2seeds: Keypoint index of chosen seeds in image I_2
        scores: Scores to rank correspondences by confidence.
                Lower scores are assumed to be more confident, consistently with Lowe's ratio scores.
                Note: scores should be between 0 and 1 for this function to work as expected.

        Returns:
            All information required for running the parallel RANSACs.
            Data is formatted so that all inputs for different RANSACs are concatenated
                along the same dimension to support different input sizes.

            im1loc: Keypoint locations in image I_1 for each RANSAC sample.
            im2loc: Keypoint locations in image I_2 for each RANSAC sample.
            ransidx: Integer identifier of the RANSAC problem.
                     This allows to distinguish inputs belonging to the same problem.
            tokp1: Index of the original keypoint in image I_1 for each RANSAC sample.
            tokp2: Index of the original keypoint in image I_2 for each RANSAC sample.
    """
    # first get an indexing representation of the assignments:
    # - ransidx holds the index of the seed for each assignment
    # - tokp1 holds the index of the keypoint in image I_1 for each assignment
    ransidx, tokp1 = torch.where(fnn_to_seed_local_consistency_map_corr)
    # - and of course tokp2 holds the index of the corresponding keypoint in image I_2
    tokp2 = fnn12[tokp1]

    # Now take the locations in the image of each considered keypoint ...
    im1abspattern = k1[tokp1]
    im2abspattern = k2[tokp2]

    # ... and subtract the location of its corresponding seed to get relative coordinates
    im1loc = im1abspattern - k1[im1seeds[ransidx]]
    im2loc = im2abspattern - k2[im2seeds[ransidx]]

    # Finally we need to sort keypoints by scores in a way that assignments to the same seed are close together
    # To achieve this we assume scores lie in (0, 1) and add the integer index of the corresponding seed
    expanded_local_scores = scores[tokp1] + ransidx.type(scores.dtype)

    sorting_perm = torch.argsort(expanded_local_scores)

    im1loc = im1loc[sorting_perm]
    im2loc = im2loc[sorting_perm]
    tokp1 = tokp1[sorting_perm]
    tokp2 = tokp2[sorting_perm]

    return im1loc, im2loc, ransidx, tokp1, tokp2

def adalam_core(k1: torch.Tensor,
                k2: torch.Tensor,
                fnn12: torch.Tensor,
                scores1: torch.Tensor,
                config: dict,
                mnn: torch.Tensor = None,
                im1shape: tuple = None,
                im2shape: tuple = None,
                o1: torch.Tensor = None,
                o2: torch.Tensor = None,
                s1: torch.Tensor = None,
                s2: torch.Tensor = None):
    """
        Call the core functionality of AdaLAM, i.e. just outlier filtering. No sanity check is performed on the inputs.

        Inputs:
            k1: keypoint locations in the source image, in pixel coordinates.
                Expected a float32 tensor with shape (num_keypoints_in_source_image, 2).
            k2: keypoint locations in the destination image, in pixel coordinates.
                Expected a float32 tensor with shape (num_keypoints_in_destination_image, 2).
            fn12: Initial set of putative matches to be filtered.
                  The current implementation assumes that these are unfiltered nearest neighbor matches,
                  so it requires this to be a list of indices a_i such that the source keypoint i is associated to the destination keypoint a_i.
                  For now to use AdaLAM on different inputs a workaround on the input format is required.
                  Expected a long tensor with shape (num_keypoints_in_source_image,).
            scores1: Confidence scores on the putative_matches. Usually holds Lowe's ratio scores.
            mnn: A mask indicating which putative matches are also mutual nearest neighbors. See documentation on 'force_seed_mnn' in the DEFAULT_CONFIG.
                 If None, it disables the mutual nearest neighbor filtering on seed point selection.
                 Expected a bool tensor with shape (num_keypoints_in_source_image,)
            im1shape: Shape of the source image. If None, it is inferred from keypoints max and min, at the cost of wasted runtime. So please provide it.
                      Expected a tuple with (width, height) or (height, width) of source image
            im2shape: Shape of the destination image. If None, it is inferred from keypoints max and min, at the cost of wasted runtime. So please provide it.
                      Expected a tuple with (width, height) or (height, width) of destination image
            o1/o2: keypoint orientations in degrees. They can be None if 'orientation_difference_threshold' in config is set to None.
                   See documentation on 'orientation_difference_threshold' in the DEFAULT_CONFIG.
                   Expected a float32 tensor with shape (num_keypoints_in_source/destination_image,)
            s1/s2: keypoint scales. They can be None if 'scale_rate_threshold' in config is set to None.
                   See documentation on 'scale_rate_threshold' in the DEFAULT_CONFIG.
                   Expected a float32 tensor with shape (num_keypoints_in_source/destination_image,)

        Returns:
            Filtered putative matches.
            A long tensor with shape (num_filtered_matches, 2) with indices of corresponding keypoints in k1 and k2.
    """
    AREA_RATIO = config['area_ratio']
    SEARCH_EXP = config['search_expansion']
    RANSAC_ITERS = config['ransac_iters']
    MIN_INLIERS = config['min_inliers']
    MIN_CONF = config['min_confidence']
    ORIENTATION_THR = config['orientation_difference_threshold']
    SCALE_RATE_THR = config['scale_rate_threshold']
    REFIT = config['refit']

    if im1shape is None:
        k1mins, _ = torch.min(k1, dim=0)
        k1maxs, _ = torch.max(k1, dim=0)
        im1shape = (k1maxs - k1mins).cpu().numpy()
    if im2shape is None:
        k2mins, _ = torch.min(k2, dim=0)
        k2maxs, _ = torch.max(k2, dim=0)
        im2shape = (k2maxs - k2mins).cpu().numpy()

    # Compute seed selection radii to be invariant to image rescaling
    R1 = np.sqrt(np.prod(im1shape[:2]) / AREA_RATIO / np.pi)
    R2 = np.sqrt(np.prod(im2shape[:2]) / AREA_RATIO / np.pi)

    # Precompute the inner distances of keypoints in image I_1
    dist1 = dist_matrix(k1, k1)

    # Select seeds
    im1seeds, im2seeds = select_seeds(dist1, R1, scores1, fnn12, mnn)

    # Find the neighboring and coherent keyopints consistent with each seed
    local_neighs_mask, rdims, im1seeds, im2seeds = extract_neighborhood_sets(
        o1, o2, s1, s2, dist1, im1seeds, im2seeds, k1, k2, R1, R2, fnn12,
        ORIENTATION_THR, SCALE_RATE_THR, SEARCH_EXP, MIN_INLIERS)

    if rdims.shape[0] == 0:
        # No seed point survived. Just output ratio-test matches. This should happen very rarely.
        absolute_im1idx = torch.where(scores1 < 0.8**2)[0]
        absolute_im2idx = fnn12[absolute_im1idx]
        return torch.stack([absolute_im1idx, absolute_im2idx], dim=1)

    # Format neighborhoods for parallel RANSACs
    im1loc, im2loc, ransidx, tokp1, tokp2 = extract_local_patterns(
        fnn12, local_neighs_mask, k1, k2, im1seeds, im2seeds, scores1)
    im1loc = im1loc / (R1 * SEARCH_EXP)
    im2loc = im2loc / (R2 * SEARCH_EXP)

    # Run the parallel confidence-based RANSACs to perform local affine verification
    inlier_idx, _, \
    inl_confidence, inlier_counts = ransac(xsamples=im1loc,
                                           ysamples=im2loc,
                                           rdims=rdims, iters=RANSAC_ITERS,
                                           refit=REFIT, config=config)

    conf = inl_confidence[ransidx[inlier_idx]]
    cnt = inlier_counts[ransidx[inlier_idx]].float()
    passed_inliers_mask = (conf >= MIN_CONF) & (cnt * (1 - 1/conf) >= MIN_INLIERS)
    accepted_inliers = inlier_idx[passed_inliers_mask]

    absolute_im1idx = tokp1[accepted_inliers]
    absolute_im2idx = tokp2[accepted_inliers]

    final_matches = torch.stack([absolute_im1idx, absolute_im2idx], dim=1)
    if final_matches.shape[0] > 1:
        return torch.unique(final_matches, dim=0)
    return final_matches


"""Simulator for bulk datasets"""

import matplotlib
import matplotlib.pyplot
import torch
import pyro
import anndata
import numpy as np
import pandas as pd
from typing import Dict, Optional
from ternadecov.stats_helpers import NegativeBinomialAltParam
from ternadecov.dataset import SingleCellDataset
from anndata import AnnData

def generate_anndata_from_sim(sim_res: Dict, sc_dataset: SingleCellDataset) -> AnnData:
    """Generate AnnData object from the simulation results

    :param sim_res: Simulation results dictonary
    :param sc_dataset: A single-cell dataset object (for the gene names)

    :return: AnnData object with simulated data
    """

    var_tmp = pd.DataFrame({"gene": list(sc_dataset.sc_anndata.var.index)})
    var_tmp = var_tmp.set_index("gene")

    return AnnData(
        X=sim_res["x_ng"].numpy(),
        var=var_tmp,
        obs=pd.DataFrame({"time": sim_res["t_m"]}),
    )

def plot_simulated_proportions(
    sim_res,
    dataset,
    show_sample_proportions=True,
    show_trajectories=True,
    figsize=(20, 10),
):
    """Plot simulated proportion results

    :param sim_res: simulation results objects
    :dataset:
    :param show_sample_proportions: show the generated proportions plot
    :param show_trajectories: show underlying trajectories plot

    :return: matplotlib axes
    """

    true_trajectories = sim_res["trajectory_params"]["trajectories_cm"]
    trajectory_type = sim_res["trajectory_params"]["type"]

    # Order the time axis
    o = torch.argsort(sim_res["t_m"])

    fig, ax = matplotlib.pyplot.subplots(
        sum((show_sample_proportions, show_trajectories)), figsize=figsize
    )

    if show_trajectories:
        # Get time range and initialize tensor
        n_samples = 1000
        t_min = torch.min(sim_res["t_m"]).item()
        t_max = torch.max(sim_res["t_m"]).item()
        t_m = np.linspace(t_min, t_max, n_samples)
        trajectories_cm = torch.zeros(dataset.num_cell_types, n_samples)

        if trajectory_type == "linear":
            a = sim_res["trajectory_params"]["a"]
            b = sim_res["trajectory_params"]["b"]
            for i in range(dataset.num_cell_types):
                trajectories_cm[i, :] = torch.Tensor(list(a[i] * x + b[i] for x in t_m))
            #trajectories_cm = torch.nn.functional.softmax(trajectories_cm, dim=0)
        elif trajectory_type == "sigmoid":
            effect_size = sim_res["trajectory_params"]["effect_size"]
            shift = sim_res["trajectory_params"]["shift"]
            for i in range(dataset.num_cell_types):
                trajectories_cm[i, :] = torch.Tensor(
                    list(sigmoid(effect_size[i] * x + shift[i]) for x in t_m)
                )
            #trajectories_cm = torch.nn.functional.softmax(trajectories_cm, dim=0)
        elif trajectory_type == "periodic":
            a = sim_res["trajectory_params"]["a"]
            b = sim_res["trajectory_params"]["b"]
            c = sim_res["trajectory_params"]["c"]
            for i in range(dataset.num_cell_types):
                trajectories_cm[i, :] = torch.Tensor(
                    list(a[i] * torch.sin(b[i] * x + c[i]) for x in t_m)
                )
            #trajectories_cm = torch.nn.functional.softmax(trajectories_cm, dim=0)
        else:
            raise ValueError(f"Unknown trajectory type {trajectory_type}")

        # Normalize trajectories_cm
        trajectories_cm = torch.nn.functional.softmax(trajectories_cm, dim=0)

        # Do plotting
        ax[0].plot(
            t_m, trajectories_cm.T,
        )
        ax[0].legend(dataset.cell_type_str_list)
        ax[0].set_title("True simulated trajectories")
        ax[0].set_xlabel("Set time")
        ax[0].set_ylabel("Proportions")

    if show_sample_proportions:
        for i in range(sim_res["cell_pop_cm"].shape[0]):
            ax[1].scatter(sim_res["t_m"][o], sim_res["cell_pop_cm"][i, o])
        ax[1].set_title("Simulated proportions in samples")
        ax[1].set_xlabel("Set time")
        ax[1].set_ylabel("Proportions")

    return ax

def simulate_data(
    w_hat_gc,
    start_time=-5,
    end_time=5,
    num_samples=100,
    lib_size_mean=1e6,
    lib_size_std=2e5,
    use_betas=False,
    dirichlet_alpha=1000,
    trajectory_type="sigmoid",
    trajectory_coef=None,
    phi_mean=0.15,
    phi_std=0.05,
    beta_mean=1.0,
    beta_std=0.1,
    trajectory_sample_params={},
    seed=None,
):
    """Simulate bulk data with compositional changes

    :param w_hat_gc: reference matrix
    :param start_time: time start
    :param end_time: time end
    :param num_samples: number of samples to simulate
    :param lib_size_mean: mean library size
    :param lib_size_std: library size standard deviation
    :param use_betas: use beta values from the reference model
    :param dirichlet_alpha: global dirichlet alpha coefficient
    :param trajectory_type: type of trajectory ('sigmoid','linear','periodic')
    :param trajectory_coef: predefined trajectory coefficients, if not provided they are sampled
    :param phi_mean: $\phi_{mean}$ value
    :param phi_std: $\phi_{std}$ values
    :param beta_mean: $\beta_{mean}$ values
    :param beta_std: $\beta_{std}$ values
    :param trajectory_sample_params: Dictionary of trajectory sample parameters
    :param seed: seed for trajectory sampling (optional)

    :return: dictionary of simulated values and underlying coefficients
    """

    num_genes = w_hat_gc.shape[0]
    num_cell_types = w_hat_gc.shape[1]

    # Get equidistant time samples
    t_m = torch.linspace(0.0, 1.0, num_samples) * (end_time - start_time) + start_time

    if trajectory_type == "sigmoid":
        proportions_sample = sample_sigmoid_proportions(
            num_cell_types=num_cell_types,
            num_samples=num_samples,
            t_m=t_m,
            dirichlet_alpha=dirichlet_alpha,
            trajectory_coefficients=trajectory_coef,
            trajectory_sample_params=trajectory_sample_params,
            seed=seed,
        )
    elif trajectory_type == "periodic":
        proportions_sample = sample_periodic_proportions(
            num_cell_types=num_cell_types,
            num_samples=num_samples,
            t_m=t_m,
            dirichlet_alpha=dirichlet_alpha,
            trajectory_coefficients=trajectory_coef,
            trajectory_sample_params=trajectory_sample_params,
            seed=seed,
        )
    elif trajectory_type == "linear":
        proportions_sample = sample_linear_proportions(
            num_cell_types,
            num_samples,
            t_m,
            dirichlet_alpha,
            trajectory_coef,
            trajectory_sample_params=trajectory_sample_params,
            seed=seed,
        )
    else:
        raise Exception("Unkown Trajectory Type")

    cell_pop_cm = proportions_sample["cell_pop_cm"]
    phi_g = (
        torch.distributions.normal.Normal(
            loc=torch.full((num_genes,), phi_mean),
            scale=torch.full((num_genes,), phi_std),
        )
        .sample()
        .abs()
    )
    beta_g = (
        torch.distributions.normal.Normal(
            loc=torch.full((num_genes,), beta_mean),
            scale=torch.full((num_genes,), beta_std),
        )
        .sample()
        .abs()
    )

    # Get celltype profiles from the model
    if use_betas:
        unnorm_w_hat_gc = w_hat_gc * beta_g[:, None]
    else:
        unnorm_w_hat_gc = w_hat_gc

    # Normalize
    w_gc = unnorm_w_hat_gc / unnorm_w_hat_gc.sum(0)

    # Sample library sizes
    lib_sizes_m = torch.normal(
        mean=torch.full([num_samples], lib_size_mean),
        std=torch.full([num_samples], lib_size_std),
    )

    mu_mg = lib_sizes_m[:, None] * torch.matmul(cell_pop_cm.T, w_gc.transpose(-1, -2))

    # Sample a full matrix using phis from main model
    x_ng = NegativeBinomialAltParam(mu=mu_mg, phi=phi_g).sample()

    return {
        "cell_pop_cm": cell_pop_cm,
        "t_m": t_m,
        "x_ng": x_ng,
        "trajectory_params": proportions_sample["trajectory_params"],
    }

def sample_trajectories(type, num_cell_types, seed = None):
    """Generate a random trajectory

    :param type: trajectory type (linear, sigmoid, periodical)
    :param num_cell_types: number of cell types in the trajectory

    :return simulated trajectories:
    """
    if type == "linear":
        return sample_linear_trajectories(num_cell_types, seed=seed)
    elif type == "periodic":
        return sample_periodic_trajectories(num_cell_types, seed=seed)
    elif type == "sigmoid":
        return sample_sigmoid_trajectories(num_cell_types, seed=seed)

######################################################
# Linear
######################################################

def sample_linear_trajectories(
    num_cell_types: int,
    seed: Optional[int] = None,
    a_min: float = 0.0,
    a_max: float = 10.0,
    b_min: float = -10.0,
    b_max: float = 10.0,
) -> Dict:
    """
    Generate a sample of linear trajectory coefficients

    :param num_cell_types: number of cell types in trajectories
    :param seed: Random seed (for reproducibility)
    :param a_min: minimum value for a coefficient
    :param a_max: maximum values for a coefficient
    :param b_min: minimum value for b coefficient
    :param b_max: maximum value for b coefficient

    :return: Dictionary coefficient and their values
    """
    if seed is not None:
        torch.manual_seed(seed)

    # y = ax+b
    a = torch.rand(num_cell_types) * (a_max - a_min) + a_min
    b = torch.rand(num_cell_types) * (b_max - b_min) + b_min

    return {"a": a, "b": b}

def sample_linear_proportions(
    num_cell_types,
    num_samples,
    t_m,
    dirichlet_alpha=1e4,
    trajectory_coef=None,
    trajectory_sample_params=None,
    seed=None,
):
    """Generate a sample of linear proportions

    :param num_cell_types: number of cell types to simulate
    :param num_samples: number of samples
    :param t_m: torch tensor of times
    :param dirichlet_alpha: multiplier for normalized dirichlet coefficients

    :return: Dictionary of coefficients
    """

    if trajectory_coef is None:
        trajectory_coef = sample_linear_trajectories(
            num_cell_types, seed=seed, **trajectory_sample_params
        )

    a = trajectory_coef["a"]
    b = trajectory_coef["b"]

    trajectories_cm = torch.zeros(num_cell_types, num_samples)
    for i in range(num_cell_types):
        trajectories_cm[i, :] = torch.Tensor(list(a[i] * x + b[i] for x in t_m))

    # Normalize trajectories_cm
    trajectories_cm = torch.nn.functional.softmax(trajectories_cm, dim=0)

    # For every sample, sample proportions from trajectory
    cell_pop_cm = torch.zeros(num_cell_types, num_samples)
    for j in range(num_samples):
        cell_pop_cm[:, j] = torch.distributions.dirichlet.Dirichlet(
            trajectories_cm[:, j] * dirichlet_alpha
        ).sample()

    return {
        "trajectory_params": {
            "type": "linear",
            "a": a,
            "b": b,
            "trajectories_cm": trajectories_cm,
        },
        "cell_pop_cm": cell_pop_cm,
    }

######################################################
# Periodic
######################################################

def sample_periodic_trajectories(
    num_cell_types: int,
    seed: Optional[int] = None,
    a_min: float = -3.0,
    a_max: float = 3.0,
    b_min: float = 0.25,
    b_max: float = 1.0,
    c_min: float = 0.0,
    c_max: float = 5.0,
) -> Dict:
    """Get a sample of coefficients for periodic trajectories

    :param num_cell_types: Number of celltypes
    :param seed: Seed for reproducibility
    :param a_min: min value for a
    :param a_max: max value for a
    :param b_min: min value for b
    :param b_max: max value for b
    :param c_min: min value for c
    :param c_max: max value for c

    :return: Dictionary of coefficients
    """

    if seed is not None:
        torch.manual_seed(seed)

    # y = a sin(b*x+c)
    a = torch.rand(num_cell_types) * (a_max - a_min) + a_min
    b = torch.rand(num_cell_types) * (b_max - b_min) + b_min
    c = torch.rand(num_cell_types) * (c_max - c_min) + c_min

    return {"a": a, "b": b, "c": c}

def sample_periodic_proportions(
    num_cell_types,
    num_samples,
    t_m,
    dirichlet_alpha=1e4,
    trajectory_coefficients=None,
    trajectory_sample_params=None,
    seed=None,
):
    """Get a sample of periodic cell proportions, optionally from a given trajectory

    :param num_cell_types: number of cell types to simulate
    :param num_samples: number of samples to simulate
    :param t_m: time points to simulate results for
    :param dirichlet_alpha: global diriechlet concentration
    :param trajectory_coefficients: trajectory, if not provided a random trajectory is drawn
    :param trajectory_sample_params: optional parameter dictionary for sampling trajectories
    :param seed: optional seed for drawing coefficients

    """
    if trajectory_coefficients is None:
        trajectory_coefficients = sample_periodic_trajectories(
            num_cell_types=num_cell_types, seed=seed, **trajectory_sample_params
        )

    a = trajectory_coefficients["a"]
    b = trajectory_coefficients["b"]
    c = trajectory_coefficients["c"]

    trajectories_cm = torch.zeros(num_cell_types, num_samples)
    for i in range(num_cell_types):
        trajectories_cm[i, :] = torch.Tensor(
            list(a[i] * torch.sin(b[i] * x + c[i]) for x in t_m)
        )

    trajectories_cm = torch.nn.functional.softmax(trajectories_cm, dim=0)

    # For every sample, sample proportions from trajectory
    cell_pop_cm = torch.zeros(num_cell_types, num_samples)
    for j in range(num_samples):
        cell_pop_cm[:, j] = torch.distributions.dirichlet.Dirichlet(
            trajectories_cm[:, j] * dirichlet_alpha
        ).sample()

    return {
        "trajectory_params": {
            "type": "periodic",
            "a": a,
            "b": b,
            "c": c,
            "trajectories_cm": trajectories_cm,
        },
        "cell_pop_cm": cell_pop_cm,
    }

######################################################
# Sigmoid
######################################################

def sigmoid(x):
    """Return sigmoid function value"""
    return 1.0 / (1.0 + np.exp(-x))

def sample_sigmoid_trajectories(
    num_cell_types,
    seed=None,
    effect_size_min=-1,
    effect_size_max=1,
    shift_min=-2,
    shift_max=2,
):
    """Return sigmoid trajectory param dictionary"""

    if seed is not None:
        torch.manual_seed(seed)

    effect_size = (
        torch.rand(num_cell_types) * (effect_size_max - effect_size_min)
        + effect_size_min
    )
    shift = torch.rand(num_cell_types) * (shift_max - shift_min) + shift_min

    return {"effect_size": effect_size, "shift": shift}

def sample_sigmoid_proportions(
    num_cell_types,
    num_samples,
    t_m,
    dirichlet_alpha=1e4,
    trajectory_coefficients=None,
    trajectory_sample_params={},
    seed=None,
):
    """Generate a sample of sigmoid proportions

    :param num_cell_types: number of cell types to simulate
    :param num_samples: number of samples
    :param t_m: torch tensor of times
    :param dirichlet_alpha: multiplier for normalized dirichlet coefficients

    :return: Dictionary of coefficients
    """
    if trajectory_coefficients is None:
        trajectory_coefficients = sample_sigmoid_trajectories(
            num_cell_types=num_cell_types, seed=seed, **trajectory_sample_params,
        )

    effect_size = trajectory_coefficients["effect_size"]
    shift = trajectory_coefficients["shift"]

    # Generate trajectories_cm
    trajectories_cm = torch.zeros(num_cell_types, num_samples)
    for i in range(num_cell_types):
        trajectories_cm[i, :] = torch.Tensor(
            list(sigmoid(effect_size[i] * x + shift[i]) for x in t_m)
        )

    trajectories_cm = torch.nn.functional.softmax(trajectories_cm, dim=0)

    # For every sample, sample proportions from trajectory
    cell_pop_cm = torch.zeros(num_cell_types, num_samples)
    for j in range(num_samples):
        cell_pop_cm[:, j] = torch.distributions.dirichlet.Dirichlet(
            trajectories_cm[:, j] * dirichlet_alpha
        ).sample()

    return {
        "trajectory_params": {
            "type": "sigmoid",
            "effect_size": effect_size,
            "shift": shift,
            "trajectories_cm": trajectories_cm,
        },
        "cell_pop_cm": cell_pop_cm,
    }

def calculate_sample_prediction_error(sim_res, pseudo_time_reg_deconv_sim) -> Dict:
    """Calculate the error at the level of individual sample proportion prediction

    :param sim_res: Simulation results to use as base truth
    :param pseudo_time_reg_deconv_sim: The trained object to simulate

    :return: Dictionary of errors
    """

    # Ground Truth
    ground_truth_cell_pop_cm = sim_res["cell_pop_cm"]
    estimated_cell_pop_cm = (
        pyro.param("cell_pop_posterior_loc_mc").clone().detach().cpu().T
    )

    l1_error = (ground_truth_cell_pop_cm - estimated_cell_pop_cm).abs().sum([0, 1])
    l1_error_norm = l1_error / estimated_cell_pop_cm.shape[-1]

    return {"l1_error": l1_error, "l1_error_norm": l1_error_norm}

def calculate_trajectory_prediction_error(
    sim_res, pseudo_time_reg_deconv_sim, n_intervals=1000
):
    """Calculate the prediction error of a deconvolution on simulated results

    :param sim_res: results of a simulation
    :param pseudo_time_reg_deconv_sim: the deconvolution object to evaluate
    :n_intervals: number of intervals over which to evaluate the results

    :return: Dictionary of results
    """

    start_time = -5
    end_time = 5
    step = (end_time - start_time) / n_intervals

    t_m = torch.arange(start_time, end_time, step)
    num_samples = t_m.shape[0]

    # Get the ground truth
    if sim_res["trajectory_params"]["type"] == "sigmoid":
        shift = sim_res["trajectory_params"]["shift"]
        effect_size = sim_res["trajectory_params"]["effect_size"]
        num_cell_types = sim_res["trajectory_params"]["effect_size"].shape[0]

        cell_pop_cm = torch.zeros(num_cell_types, num_samples)
        for i in range(num_cell_types):
            cell_pop_cm[i, :] = torch.Tensor(
                list(sigmoid(effect_size[i] * x + shift[i]) for x in t_m)
            )
        ground_truth_proportions_cm = torch.nn.functional.softmax(cell_pop_cm, dim=0)
    elif sim_res["trajectory_params"]["type"] == "linear":
        a = sim_res["trajectory_params"]["a"]
        b = sim_res["trajectory_params"]["b"]
        num_cell_types = sim_res["trajectory_params"]["trajectories_cm"].shape[0]

        cell_pop_cm = torch.zeros(num_cell_types, num_samples)
        for i in range(num_cell_types):
            cell_pop_cm[i, :] = torch.Tensor(list(a[i] * x + b[i] for x in t_m))
        ground_truth_proportions_cm = torch.nn.functional.softmax(cell_pop_cm, dim=0)
    elif sim_res["trajectory_params"]["type"] == "periodic":
        a = sim_res["trajectory_params"]["a"]
        b = sim_res["trajectory_params"]["b"]
        c = sim_res["trajectory_params"]["c"]
        num_cell_types = sim_res["trajectory_params"]["trajectories_cm"].shape[0]

        cell_pop_cm = torch.zeros(num_cell_types, num_samples)
        for i in range(num_cell_types):
            cell_pop_cm[i, :] = torch.Tensor(
                list(a[i] * torch.sin(b[i] * x + c[i]) for x in t_m)
            )
        ground_truth_proportions_cm = torch.nn.functional.softmax(cell_pop_cm, dim=0)
    else:
        raise Exception(
            f'Unknown trajectory type { sim_res["trajectory_params"]["type"] }'
        )

    # Get the predictions
    traj = pseudo_time_reg_deconv_sim.population_proportion_model.get_composition_trajectories(
        dataset=pseudo_time_reg_deconv_sim.dataset, n_intervals=n_intervals
    )
    ret_vals = traj

    predicted_composition_cm = ret_vals["norm_comp_tc"].T

    # Calculate L1 and L2 losses
    L1_error = (
        (ground_truth_proportions_cm - predicted_composition_cm).abs().sum([0, 1])
    )
    L1_error_norm = L1_error / n_intervals
    L2_error = (
        (ground_truth_proportions_cm - predicted_composition_cm)
        .pow(2)
        .sum([0, 1])
        .sqrt()
    )
    L2_error_norm = L2_error / n_intervals

    # Calculate L1 and L2 losses on the trajectory shapes

    # Normalize by cell type summing to 1
    ground_truth_proportions_norm_cm = (
        ground_truth_proportions_cm / ground_truth_proportions_cm.sum(-2)
    )
    predicted_composition_norm_cm = (
        predicted_composition_cm / predicted_composition_cm.sum(-2)
    )

    shape_L1_error = (
        (ground_truth_proportions_norm_cm - predicted_composition_norm_cm)
        .abs()
        .sum([0, 1])
    )

    return {
        "L1_error": L1_error,
        "L1_error_norm": L1_error_norm,
        "L2_error": L2_error,
        "L2_error_norm": L2_error_norm,
        "shape_L1_error": shape_L1_error,
    }


from pydca.contact_visualizer.contact_visualizer import DCAVisualizer
from pydca.meanfield_dca.meanfield_dca import MeanFieldDCA
from pydca.plmdca.plmdca import PlmDCA
from pydca.fasta_reader import fasta_reader
from pydca.sequence_backmapper import scoring_matrix
from pydca.plmdca.plmdca import PlmDCA as PseudoLikelihoodMaxDCA
from pydca.main import configure_logging
from Bio import pairwise2
import matplotlib.pyplot as plt
from scipy.optimize import minimize as scipy_minimize
from  .convolution import Convolution
from .inputreader import InputReader
from .cmdargs import CmdArgs
import subprocess
import numpy as np
import logging
import os, errno
import glob
from datetime import datetime
import pickle
import random
from pathlib import Path
from argparse import ArgumentParser
import sys

logger = logging.getLogger(__name__)

class CocoNetException(Exception):
    """Raise exceptions related to CocoNet computation.
    """

class CoCoNet:
    """Implements RNA contact prediction using direct coupling analysis enhanced by
    a simple convolutional neural network.
    """

    def __init__(self, data_dir, linear_dist=None, contact_dist=None):
        """Initializes CocoNet instance.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            dir_msa_files : str
                Path to directory containing the MSA files.
            dir_pdb_file : str
                Path to the directory containing the PDB files.
            dir_refseq_files : str
                Path to the directory containing reference sequence files.
            linear_dist : int
                Distance between sites in reference sequence
            contact_dist : float
                Maximum distance between two residues in PDB file to be considered
                contacts.
        """
        self.__data_dir = os.path.abspath(data_dir)
        self.__linear_dist = linear_dist if linear_dist is not None else 4
        self.__contact_dist = contact_dist if contact_dist is not None else 10.0

        pdb_chains_list_file = os.path.join(self.__data_dir, 'CCNListOfPDBChains.txt')
        msa_files_list_file = os.path.join(self.__data_dir, 'CCNListOfMSAFiles.txt')
        pdb_files_list_file = os.path.join(self.__data_dir, 'CCNListOfPDBFiles.txt')
        input_reader = InputReader()
        self.__msa_file_names_list = input_reader.read_from_one_column_text_file(msa_files_list_file)
        self.__pdb_chains_list = input_reader.read_from_one_column_text_file(pdb_chains_list_file)
        self.__pdb_file_names_list = input_reader.read_from_one_column_text_file(pdb_files_list_file)
        self.__msa_files_dir = os.path.join(self.__data_dir, 'MSA')
        self.__refseqs_dir = os.path.join(self.__data_dir, 'sequences')
        self.__pdb_files_dir = os.path.join(self.__data_dir, 'PDBFiles')
        self.__secstruct_files_dir = os.path.join(self.__data_dir, 'secstruct')
        self.__msa_files_list  = [
            os.path.abspath(os.path.join(self.__msa_files_dir, msa_file + '.faclean')) for msa_file in self.__msa_file_names_list
        ]

        self.__refseqs = self.get_refseqs()
        self.__refseqs_len = self.get_refseqs_len()

        logmsg  = """
            Data directory          : {},
            PDB chains list file    : {},
            MSA files list file     : {},
            PDB files list file     : {},
        """.format(self.__data_dir, pdb_chains_list_file,
            msa_files_list_file, pdb_files_list_file,
        )

        logger.info(logmsg)
        return None

    @property
    def pdb_file_names_list(self):
        return self.__pdb_file_names_list

    @property
    def msa_file_names_list(self):
        return self.__msa_file_names_list

    @property
    def pdb_chains_list(self):
        return self.__pdb_chains_list

    def map_pdb_id_to_family(self):
        """Mapps PDB ID to family name.

        Parameters
        ----------
            self : CocoNet(self, data_dir, linear_dist=None, contact_dist=None)

        Returns
        -------
            pdb_id_to_fam_name : dict
                pdb_id_to_fam_name[pdb_id]=fam_name
        """
        pdb_id_to_fam_name = dict()
        for pdb_id, fam_name in zip(self.__pdb_file_names_list, self.__msa_file_names_list):
            pdb_id_to_fam_name[pdb_id] = fam_name
        return pdb_id_to_fam_name

    @staticmethod
    def _to_dict(files_list):
        """Puts a list of file paths into a dictionary

        Parameters
        ----------
            files_list : list
                A list of file paths

        Returns
        -------
            files_dict : dict
                A dictionary whose keys are basenames of files and values file path.
        """
        files_dict = dict()
        for f in files_list:
            basename, _ = os.path.splitext(os.path.basename(f))
            files_dict[basename] = f
        return files_dict

    def get_refseq_files_list(self):
        """
        """
        refseq_files_list = [
            os.path.join(self.__refseqs_dir, pdb_file[:4] + '.fa') for pdb_file in self.__pdb_file_names_list
        ]
        return tuple(refseq_files_list)

    def get_pdb_files_list(self):
        """
        """
        pdb_files_list = [
            os.path.join(self.__pdb_files_dir, pdb_file + '.pdb') for pdb_file in self.__pdb_file_names_list
        ]
        return tuple(pdb_files_list)

    def create_directories(self, dir_path):
        """Creates (nested) directory given path.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            dir_path : str
                Directory path.

        Returns
        -------
            None : None
        """

        try:
            os.makedirs(dir_path)
        except OSError as e:
            if e.errno !=errno.EEXIST:
                logger.error('Unable to create directory using path {}'.format(
                    dir_path)
                )
                raise
        return None

    def recompute_dca_data(self, msa_files_list=None, pickled_data=None):
        """Checks the last modification time of MSA files and pickled DCA data.
        If any of the MSA files are modified recently compared to pickled DCA
        data or pickled data does not exist, return values is True else False.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            msa_files_list : list
                A list of RNA MSA files.
            pickled_data : str
                Pickled DCA data

        Returns
        -------
            True or False : bool
        """

        if any([msa_files_list, pickled_data]) is None:
            logger.error('\n\tYou need to supply all keyword arguments')
            raise CocoNetException
        if not os.path.exists(pickled_data): return True
        mtime_msa_files = max([os.path.getmtime(f) for f in msa_files_list])
        mtime_pickled_data = os.path.getmtime(pickled_data)
        if mtime_msa_files > mtime_pickled_data:
            return True
        return False

    def compute_mfdca_DI_scores(self):
        """Computes the mean-field DCA score of all the RNA families in the directory
        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class

        Returns
        -------
            all_dca_data : dict
        """
        all_dca_data = dict()
        for msa_file in self.__msa_files_list:
            logger.info('\n\tMSA file: {}'.format(msa_file))
            famname, _ext = os.path.splitext(os.path.basename(msa_file))
            logger.info('\n\tFamily name: {}'.format(famname))
            try:
                mfdca = MeanFieldDCA(msa_file, 'rna')
                dca_data = mfdca.compute_sorted_DI_APC()
            except Exception:
                raise
            else:
                all_dca_data[famname] = dca_data
        # The DCA scores are in a list of tuples. Lets convert them to a dict
        all_dca_data_dict = dict()
        for rna_fam in all_dca_data:
            all_dca_data_dict[rna_fam] = dict((pair, score) for pair, score in all_dca_data[rna_fam])
        return all_dca_data_dict

    def compute_plmdca_FN_APC_scores(self, max_iterations=500000, num_threads=1, verbose=False):
        """
        """
        all_dca_data = dict()
        for msa_file in self.__msa_files_list:
            logger.info('\n\tMSA file: {}'.format(msa_file))
            famname, _ext = os.path.splitext(os.path.basename(msa_file))
            logger.info('\n\tFamily name: {}'.format(famname))
            try:
                plmdca_inst = PseudoLikelihoodMaxDCA(msa_file, 'rna', max_iterations=max_iterations, verbose=verbose, num_threads=num_threads)
                dca_data = plmdca_inst.compute_sorted_FN_APC()
            except Exception:
                raise
            else:
                all_dca_data[famname] = dca_data
        # The DCA scores are in a list of tuples. Lets convert them to a dict
        all_dca_data_dict = dict()
        for rna_fam in all_dca_data:
            all_dca_data_dict[rna_fam] = dict((pair, score) for pair, score in all_dca_data[rna_fam])
        return all_dca_data_dict

    def get_pdb_data(self):
        """Computes mapped PDB contacts for multiple RNA families. The computed
        mapped PDB data is pickled and only recomputed if any of reference sequence
        files, PDB files or PDB chain metadata file is updated.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class

        Returns
        -------
            mapped_pdb_data : dict
                A dictionary whose keys are RNA familiy names and values dictionaries
                that have site pair keys and PDB data values.
        """
        refseq_files_list = self.get_refseq_files_list()
        pdb_files_list = self.get_pdb_files_list()

        logger.info('\n\tObtaining mapped PDB data')
        txtfreader = InputReader()
        mapped_pdb_data = dict()
        for chain_id, pdb_file, refseq_file, msa_file in zip(self.__pdb_chains_list, pdb_files_list, refseq_files_list, self.__msa_file_names_list):
            curr_pdb_data, _missing, _refseq_len = txtfreader.get_mapped_pdb_data(pdb_chain_id=chain_id,
                refseq_file=refseq_file, pdb_file=pdb_file, linear_dist=self.__linear_dist,
                contact_dist=self.__contact_dist
            )
            # self.__msa_file_names_list  contains the list of MSA files, not the full path of the files
            famname, _ext =  os.path.splitext(msa_file)
            mapped_pdb_data[famname] = curr_pdb_data
        return mapped_pdb_data

    def get_refseqs(self):
        """Obtains reference sequences of several RNA famlies from fasta formatted file.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class

        Returns
        -------
            reference_sequenes_dict : dict()
            reference_sequences_dict[FAMILY_NAME] = sequence
        """

        refseq_files_list = self.get_refseq_files_list()
        logger.info('\n\tObtaining reference sequences from FASTA files')
        reference_sequences_dict = dict()
        for refseq_file, msa_file_basename in zip(refseq_files_list, self.__msa_file_names_list):
            # if reference sequence file contains multiple sequences, take the first one.
            reference_sequences_dict[msa_file_basename.strip()] = fasta_reader.get_alignment_from_fasta_file(refseq_file)[0].strip()
        return reference_sequences_dict

    def get_refseqs_len(self):
        """Obtains length of reference sequence for each RNA family

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class

        Returns
        -------
            refseqs_len_dict : dict()
                refseqs_len_dict[FAMILY_NAME] = refseq_length
        """
        refseqs_dict = self.__refseqs
        logger.info('\n\tObtaining length of reference sequences for {} RNA families'.format(len(refseqs_dict)))

        refseqs_len_dict = {
            fam : len(refseqs_dict[fam]) for fam in refseqs_dict
        }
        return refseqs_len_dict

    def objective_function(self, weight_matrix, dca_data_train, pdb_data_train):
        """Computes the total objective function (for the entire training data set)

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            weight_matrix : np.array
                A 1d numpy array of weights
            dca_data_train : dict
                DCA data for training.
            pdb_data_train : dict
                PDB data for training.

        Returns
        -------
            total_cost : float
                Total value of the objective/error function.
        """
        fm_dim_size = int(np.sqrt(weight_matrix.size))
        assert fm_dim_size * fm_dim_size == weight_matrix.size
        conv_inst = Convolution(fm_dim_size)
        total_cost = 0.0
        logger.info('\n\tComputing objective function using filter matrix of size: {}'.format(weight_matrix.size))
        for fam in dca_data_train:
            fam_dca_scores = dca_data_train[fam]
            fam_pdb_contacts = pdb_data_train[fam]
            fam_refseq_len = self.__refseqs_len[fam]
            fam_cost = conv_inst.objective_function(fam_pdb_contacts, fam_dca_scores, weight_matrix, fam_refseq_len)
            total_cost += fam_cost
        return total_cost

    def objective_function_WC_and_NONWC_pairs(self, weight_matrix, dca_data_train, pdb_data_train):
        """Computes value of cost function when contacts are categorized as WC and non-WC nucleotide
        pairs.

        Parameters
        ----------
            self :CocoNet
                CocoNet(self, data_dir, linear_dist=None, contact_dist=None)
            weight_matrix : np.array()
                1d numpy array of weights. Its size must be twice the size of
                the filter matrix used so as to accommodate WC and non-WC weights
                together.
            dca_data_train : dict()
                DCA data for training set families.
            pdb_data_train : dict()
                PDB data for training set families.

        Returns
        -------
            total_cost : float
                value of cost function for at a particular iteration (value of weight matrix).
        """

        fm_dim_size = int(np.sqrt((weight_matrix.size/2)))
        assert weight_matrix.size == 2*fm_dim_size * fm_dim_size
        conv_inst = Convolution(fm_dim_size)
        total_cost = 0.0
        logger.info('\n\tComputing objective function for WC and non-WC residue pairs using filter matrix of size: {}'.format(weight_matrix.size))
        for fam in dca_data_train:
            fam_dca_scores = dca_data_train[fam]
            fam_pdb_contacts = pdb_data_train[fam]
            fam_refseq = self.__refseqs[fam]
            fam_cost = conv_inst.objective_function_WC_and_NONWC_pairs(fam_pdb_contacts, fam_dca_scores, weight_matrix, fam_refseq)
            total_cost += fam_cost
        return total_cost

    def total_gradients(self, weight_matrix, dca_data_train, pdb_data_train):
        """Computes the total gradient of the error function. That is, for more than
        one RNA family.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            weight_matrix : np.array
                A 1d numpy array of weights
            dca_data_train : dict
                DCA data for training.
            pdb_data_train : dict
                PDB data for training.

        Returns
        -------
        total_gradients : np.array
            A 1d numpy array of the total gradient

        """
        logger.info('\n\tFilter Matrix\n{}'.format(weight_matrix))
        fm_dim_size = int(np.sqrt(weight_matrix.size))
        assert fm_dim_size * fm_dim_size == weight_matrix.size
        conv_inst = Convolution(fm_dim_size)
        total_gradients = 0.0 # will be promoted to numpy array when added to another array
        for fam in dca_data_train:
            fam_dca_scores = dca_data_train[fam]
            fam_pdb_contacts = pdb_data_train[fam]
            fam_refseq_len = self.__refseqs_len[fam]
            fam_gradient = conv_inst.gradients(fam_pdb_contacts, fam_dca_scores, weight_matrix, fam_refseq_len)
            total_gradients += fam_gradient
        return total_gradients

    def total_gradients_WC_and_NONWC_pairs(self, weight_matrix, dca_data_train, pdb_data_train):
        """Computes the total gradient of the error function for WC and non-WC contact
        pair classification for test RNA families.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            weight_matrix : np.array
                A 1d numpy array of weights
            dca_data_train : dict
                DCA data for training.
            pdb_data_train : dict
                PDB data for training.

        Returns
        -------
        total_gradients : np.array
            A 1d numpy array of the total gradient

        """

        logger.info('\n\tFilter Matrix\n{}'.format(weight_matrix))
        fm_dim_size = int(np.sqrt(weight_matrix.size/2))
        assert weight_matrix.size == 2*fm_dim_size * fm_dim_size
        conv_inst = Convolution(fm_dim_size)
        total_gradients = 0.0 # will be promoted to numpy array when added to another array
        for fam in dca_data_train:
            fam_dca_scores = dca_data_train[fam]
            fam_pdb_contacts = pdb_data_train[fam]
            fam_refseq = self.__refseqs[fam]
            fam_gradient = conv_inst.gradients_WC_and_NONWC_pairs(fam_pdb_contacts, fam_dca_scores, weight_matrix, fam_refseq)
            total_gradients += fam_gradient
        return total_gradients

    def train(self, initial_weight_matrix, dca_data_train, pdb_data_train):
        """Performs gradient decent using scipy.optimize.minimize

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            initial_weight_matrix : np.array
                A 1d numpy array of initial weights.
            dca_data_train : dict
                DCA data for training.
            pdb_data_train : dict
                PDB data for training.

        Returns
        -------
            result_lbfgs :  scipy.optimize.optimize.OptimizeResult
                A instance of scipy.optimize.optimize.OptimizeResult class. It has
                the following attributes:
                result_lbfgs.fun  : float
                    Value of minimized  objective function
                result_lbfgs.jac : np.array
                    A 1d numpy array of the hessian inverse
                result_lbfgs.nit : int
                    The number of minimization iterations.
                result_lbfgs.success : bool
                    Success status of minimization
                result_lbfgs.x : np.array
                    A 1d numpy array of optimized weights.
                Note: there are more attributes that can be displayed using, for
                example, print(result_lbfgs_b)
        """
        result_lbfgs = scipy_minimize(self.objective_function, initial_weight_matrix,
            method='L-BFGS-B', jac=self.total_gradients,
            args=(dca_data_train, pdb_data_train)
        )
        if not result_lbfgs.success: raise CocoNetException('Iteration not converged')
        return result_lbfgs

    def train_3x3(self, dca_data_train, pdb_data_train):
        """Computes a 3x3 trained filter matrix.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            dca_data_train : dict
                Training DCA data.
            pdb_data_train : dict
                Training PDB data.

        Returns
        -------
            lbfgs_result.x : np.array
                A 1d numpy array of trained filter matrix elements.
        """
        weight_matrix = np.zeros(shape=(9, ), dtype=np.float64)

        lbfgs_result = self.train(weight_matrix, dca_data_train, pdb_data_train)
        return lbfgs_result.x

    def train_5x5(self, dca_data_train, pdb_data_train):
        """Computes a 5x5 trained filter matrix.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            dca_data_train : dict
                Training DCA data.
            pdb_data_train : dict
                Training PDB data.

        Returns
        -------
            lbfgs_result.x : np.array
                A 1d numpy array of trained filter matrix elements.
        """
        weight_matrix = np.zeros(shape=(25, ), dtype=np.float64)
        lbfgs_result = self.train(weight_matrix, dca_data_train, pdb_data_train)
        return lbfgs_result.x

    def train_7x7(self, dca_data_train, pdb_data_train):
        """Computes a 7x7 trained filter matrix.

        Parameters
        ----------
            self : CocoNet
                An instance of CocoNet class.
            dca_data_train : dict
                Training DCA data.
            pdb_data_train : dict
                Training DCA data.

        Returns
        -------
            lbfgs_results.x : np.array
                A 1d numpy array of trained filter matrix elements.
        """
        weight_matrix = np.zeros(shape=(49, ), dtype=np.float64)
        lbfgs_result = self.train(weight_matrix, dca_data_train, pdb_data_train)
        return lbfgs_result.x

    def train_WC_and_NONWC(self, initial_weight_matrix, dca_data_train, pdb_data_train):
        """
        """
        result_lbfgs = scipy_minimize(self.objective_function_WC_and_NONWC_pairs,
            initial_weight_matrix, method='L-BFGS-B',
            jac=self.total_gradients_WC_and_NONWC_pairs,
            args=(dca_data_train, pdb_data_train)
        )
        if not result_lbfgs.success: raise CocoNetException('Iteration not converged')
        return result_lbfgs

    def train_WC_and_NONWC_3x3(self, dca_data_train, pdb_data_train):
        """
        """
        weight_matrix = np.zeros(shape=(18, ), dtype=np.float64)
        lbfgs_result = self.train_WC_and_NONWC(weight_matrix, dca_data_train, pdb_data_train)
        return lbfgs_result.x

    def train_WC_and_NONWC_5X5(self, dca_data_train, pdb_data_train):
        """
        """
        weight_matrix = np.zeros(shape=(50, ), dtype=np.float64)
        lbfgs_result =  self.train_WC_and_NONWC(weight_matrix, dca_data_train, pdb_data_train)
        return lbfgs_result.x

    def train_WC_and_NONWC_7x7(self, dca_data_train, pdb_data_train):
        """
        """
        weight_matrix = np.zeros(shape=(98, ), dtype=np.float64)
        lbfgs_result = self.train_WC_and_NONWC(weight_matrix, dca_data_train, pdb_data_train)
        return lbfgs_result.x

    def cross_validation_single_matrix(self, matrix_size=None, wc_and_nwc=False, num_batches=5, output_dir=None, on_plm=False,
            verbose=False, num_threads=None, max_iterations=None, num_trials=1):
        """Performs cross validation of CocoNet
        """

        pdb_data = self.get_pdb_data()
        if on_plm:
            dca_data=  self.compute_plmdca_FN_APC_scores(max_iterations=max_iterations, num_threads=num_threads, verbose=verbose)
            if wc_and_nwc and output_dir is None:
                output_dir = f'CoCoNet_plmDCA_CrossValidation_Output_2x{matrix_size}x{matrix_size}-' + datetime.now().strftime('%Y-%m-%d-%H_%M_%S')
            if wc_and_nwc is False and output_dir is None:
                output_dir = f'CoCoNet_plmDCA_CrossValidation_Output_{matrix_size}x{matrix_size}-' + datetime.now().strftime('%Y-%m-%d-%H_%M_%S')

        else: # uses mean-field DCA
            dca_data = self.compute_mfdca_DI_scores()
            if wc_and_nwc and output_dir is None:
                output_dir = f'CoCoNet_mfDCA_CrossValidation_Output_2x{matrix_size}x{matrix_size}-' + datetime.now().strftime('%Y-%m-%d-%H_%M_%S')
            if wc_and_nwc is False and output_dir is None:
                output_dir = f'CoCoNet_mfDCA_CrossValidation_Output_{matrix_size}x{matrix_size}-' + datetime.now().strftime('%Y-%m-%d-%H_%M_%S')

        fams_in_DCA = list(dca_data.keys())
        fams_in_PDB = list(pdb_data.keys())
        for fam in fams_in_DCA: assert fam in fams_in_PDB
        batch_len = len(fams_in_PDB)//num_batches

        for i in range(num_trials):
            # create output destination directories

            trial_dir = 'trial_{}'.format(i + 1)
            trial_output_dir =  os.path.join(output_dir, trial_dir)
            # shuffle the list of RNAs
            random.shuffle(fams_in_PDB)
            #divide families into batches
            for j in range(num_batches):
                trial_batch_output_dir = os.path.join(trial_output_dir, 'fold_{}'.format(j + 1))
                self.create_directories(trial_batch_output_dir)
                lower_bound = j * batch_len
                upper_bound = lower_bound + batch_len
                batch_j = fams_in_PDB[lower_bound:upper_bound] if j < (num_batches - 1) else fams_in_PDB[lower_bound:]
                # take batch_j as a test set
                testset_fams = batch_j
                training_fams = [fam for fam in fams_in_PDB if fam not in testset_fams]
                dca_data_train_j = { fam : dca_data[fam] for fam in training_fams }
                pdb_data_train_j = {fam : pdb_data[fam] for fam in training_fams}

                metadata_file = os.path.join(trial_batch_output_dir, 'metadata_fold_{}.txt'.format(j + 1))

                with open(metadata_file, 'w') as fh:
                    fh.write('Testset RNA Families \n')
                    for counter, fam in enumerate(testset_fams, start=1): fh.write('{}\t{}\n'.format(counter, fam))
                    fh.write('Training RNA Families\n')
                    for counter, fam in enumerate(training_fams, start=1): fh.write('{}\t{}\n'.format(counter, fam))
                # perform training
                base_header = 'Coconet cross validation result for {} filter matrix.\nTotal number of training families: {}'
                # 3x3
                if matrix_size == 3 and not wc_and_nwc:
                    mat_3x3 = self.train_3x3(dca_data_train_j, pdb_data_train_j)
                    outfile_3x3 = os.path.join(trial_batch_output_dir, 'params_3x3.txt')
                    header_mat_3x3 = base_header.format('3x3', len(training_fams))
                    np.savetxt(outfile_3x3, mat_3x3, header=header_mat_3x3)
                if matrix_size == 3 and wc_and_nwc:
                    mat_WCNWC_3x3 = self.train_WC_and_NONWC_3x3(dca_data_train_j, pdb_data_train_j)
                    outfile_WCNWC_3x3 = os.path.join(trial_batch_output_dir, 'params_WCNWC_3x3.txt')
                    header_mat_WCNWC_3x3 = base_header.format('WCNWC 3x3', len(training_fams))
                    np.savetxt(outfile_WCNWC_3x3, mat_WCNWC_3x3, header=header_mat_WCNWC_3x3)
                # 5x5
                if matrix_size == 5 and not wc_and_nwc:
                    mat_5x5 = self.train_5x5(dca_data_train_j, pdb_data_train_j)
                    outfile_5x5 = os.path.join(trial_batch_output_dir, 'params_5x5.txt')
                    header_mat_5x5 = base_header.format('5x5', len(training_fams))
                    np.savetxt(outfile_5x5, mat_5x5, header=header_mat_5x5)
                if matrix_size == 5 and wc_and_nwc:
                    mat_WCNWC_5x5 = self.train_WC_and_NONWC_5X5(dca_data_train_j, pdb_data_train_j)
                    outfile_WCNWC_5x5 = os.path.join(trial_batch_output_dir, 'params_WCNWC_5x5.txt')
                    header_mat_WCNWC_5x5 = base_header.format('WCNWC 5x5', len(training_fams))
                    np.savetxt(outfile_WCNWC_5x5, mat_WCNWC_5x5, header=header_mat_WCNWC_5x5)
                # 7x7
                if matrix_size == 7 and not wc_and_nwc:
                    mat_7x7 = self.train_7x7(dca_data_train_j, pdb_data_train_j)
                    outfile_7x7 = os.path.join(trial_batch_output_dir, 'params_7x7.txt')
                    header_mat_7x7 = base_header.format('7x7', len(training_fams))
                    np.savetxt(outfile_7x7, mat_7x7, header=header_mat_7x7)
                if matrix_size == 7 and wc_and_nwc:
                    mat_WCNWC_7x7 = self.train_WC_and_NONWC_7x7(dca_data_train_j, pdb_data_train_j)
                    outfile_WCNWC_7x7 = os.path.join(trial_batch_output_dir, 'params_WCNWC_7x7.txt')
                    header_mat_WCNWC_7x7 = base_header.format('WCNWC 7x7', len(training_fams))
                    np.savetxt(outfile_WCNWC_7x7, mat_WCNWC_7x7, header=header_mat_WCNWC_7x7)
        return None

    def cross_validation_all_matrices(self, num_batches=5, ouput_dir=None, on_plm=False, num_threads=1,
            num_trials=1, max_iterations=None, verbose=False):
        """
        """
        logger.info('\n\tTraining CoCoNet for all filter matrices, i.e., for 3x3, 2x3x3, 5x5, 2x5x5, 7x7 and 2x7x7')

        pdb_data = self.get_pdb_data()
        if on_plm:
            dca_data=  self.compute_plmdca_FN_APC_scores(max_iterations=max_iterations, num_threads=num_threads, verbose=verbose)
            output_dir = f'CoCoNet_plmDCA_CrossValidation_Output_All_Matrices-' + datetime.now().strftime('%Y-%m-%d-%H_%M_%S')

        else: # uses mean-field DCA
            dca_data = self.compute_mfdca_DI_scores()
            output_dir = f'CoCoNet_mfDCA_CrossValidation_Output_All_Matrices-' + datetime.now().strftime('%Y-%m-%d-%H_%M_%S')

        fams_in_DCA = list(dca_data.keys())
        fams_in_PDB = list(pdb_data.keys())
        for fam in fams_in_DCA: assert fam in fams_in_PDB
        batch_len = len(fams_in_PDB)//num_batches

        for i in range(num_trials):
            # create output destination directories

            trial_dir = 'trial_{}'.format(i + 1)
            trial_output_dir =  os.path.join(output_dir, trial_dir)
            # shuffle the list of RNAs
            random.shuffle(fams_in_PDB)
            #divide families into batches
            for j in range(num_batches):
                trial_batch_output_dir = os.path.join(trial_output_dir, 'fold_{}'.format(j + 1))
                self.create_directories(trial_batch_output_dir)
                lower_bound = j * batch_len
                upper_bound = lower_bound + batch_len
                batch_j = fams_in_PDB[lower_bound:upper_bound] if j < (num_batches - 1) else fams_in_PDB[lower_bound:]
                # take batch_j as a test set
                testset_fams = batch_j
                training_fams = [fam for fam in fams_in_PDB if fam not in testset_fams]
                dca_data_train_j = { fam : dca_data[fam] for fam in training_fams }
                pdb_data_train_j = {fam : pdb_data[fam] for fam in training_fams}

                metadata_file = os.path.join(trial_batch_output_dir, 'metadata_fold_{}.txt'.format(j + 1))

                with open(metadata_file, 'w') as fh:
                    fh.write('Testset RNA Families \n')
                    for counter, fam in enumerate(testset_fams, start=1): fh.write('{}\t{}\n'.format(counter, fam))
                    fh.write('Training RNA Families\n')
                    for counter, fam in enumerate(training_fams, start=1): fh.write('{}\t{}\n'.format(counter, fam))
                # perform training
                base_header = 'Coconet cross validation result for {} filter matrix.\nTotal number of training families: {}'
                # 3x3
                mat_3x3 = self.train_3x3(dca_data_train_j, pdb_data_train_j)
                outfile_3x3 = os.path.join(trial_batch_output_dir, 'params_3x3.txt')
                header_mat_3x3 = base_header.format('3x3', len(training_fams))
                np.savetxt(outfile_3x3, mat_3x3, header=header_mat_3x3)

                # 2x3x3
                mat_WCNWC_3x3 = self.train_WC_and_NONWC_3x3(dca_data_train_j, pdb_data_train_j)
                outfile_WCNWC_3x3 = os.path.join(trial_batch_output_dir, 'params_WCNWC_3x3.txt')
                header_mat_WCNWC_3x3 = base_header.format('WCNWC 3x3', len(training_fams))
                np.savetxt(outfile_WCNWC_3x3, mat_WCNWC_3x3, header=header_mat_WCNWC_3x3)
                # 5x5
                mat_5x5 = self.train_5x5(dca_data_train_j, pdb_data_train_j)
                outfile_5x5 = os.path.join(trial_batch_output_dir, 'params_5x5.txt')
                header_mat_5x5 = base_header.format('5x5', len(training_fams))
                np.savetxt(outfile_5x5, mat_5x5, header=header_mat_5x5)

                # 2x5x5
                mat_WCNWC_5x5 = self.train_WC_and_NONWC_5X5(dca_data_train_j, pdb_data_train_j)
                outfile_WCNWC_5x5 = os.path.join(trial_batch_output_dir, 'params_WCNWC_5x5.txt')
                header_mat_WCNWC_5x5 = base_header.format('WCNWC 5x5', len(training_fams))
                np.savetxt(outfile_WCNWC_5x5, mat_WCNWC_5x5, header=header_mat_WCNWC_5x5)
                # 7x7
                mat_7x7 = self.train_7x7(dca_data_train_j, pdb_data_train_j)
                outfile_7x7 = os.path.join(trial_batch_output_dir, 'params_7x7.txt')
                header_mat_7x7 = base_header.format('7x7', len(training_fams))
                np.savetxt(outfile_7x7, mat_7x7, header=header_mat_7x7)

                # 2x7x7
                mat_WCNWC_7x7 = self.train_WC_and_NONWC_7x7(dca_data_train_j, pdb_data_train_j)
                outfile_WCNWC_7x7 = os.path.join(trial_batch_output_dir, 'params_WCNWC_7x7.txt')
                header_mat_WCNWC_7x7 = base_header.format('WCNWC 7x7', len(training_fams))
                np.savetxt(outfile_WCNWC_7x7, mat_WCNWC_7x7, header=header_mat_WCNWC_7x7)

        return None

# end of class CoCoNet

def execute_from_command_line(matrix_size=None, wc_and_nwc=False, num_trials=1,
        on_plm=False, verbose=False, output_dir=None, max_iterations=None, num_threads=None, for_all_matrices=False):
    """
    """
    if matrix_size is None: matrix_size = 3 # use this default values to annotate do ouput
    #directory names consistent with the default values in argparser.

    if verbose: configure_logging()
    logger.info('\n\tTraining CoCoNet')
    dataset_dir = Path(__file__).parent.parent / 'RNA_DATASET'
    coconet_inst = CoCoNet(dataset_dir)
    if for_all_matrices:
        coconet_inst.cross_validation_all_matrices(on_plm = on_plm, max_iterations = max_iterations, num_threads = num_threads,
            verbose = verbose, num_trials = num_trials
        )
    else:
        coconet_inst.cross_validation_single_matrix(matrix_size, num_threads = num_threads, wc_and_nwc = wc_and_nwc,
            on_plm = on_plm, verbose = verbose, max_iterations = max_iterations, num_trials = num_trials
        )

    return None

def train_coconet():
    """
    """
    parser = ArgumentParser()

    # This argument is added to help run help message when no positional argument is supplied
    parser.add_argument('run', help='Execute CoCoNet training')
    parser.add_argument(CmdArgs.verbose_optional, help=CmdArgs.verbose_optional_help, action='store_true')
    parser.add_argument(CmdArgs.matrix_size, help=CmdArgs.matrix_size_help,  type=int, choices=(3, 5, 7), default=3)
    parser.add_argument(CmdArgs.wc_and_nwc_optional, help=CmdArgs.wc_and_nwc_optional_help, action='store_true')
    parser.add_argument(CmdArgs.max_iterations_optional, help=CmdArgs.max_iterations_help, type=int, default=500000)
    parser.add_argument(CmdArgs.num_threads_optional, help=CmdArgs.num_threads_help, type=int, default=1)
    parser.add_argument(CmdArgs.on_plm_optional, help=CmdArgs.on_plm_optional_help, action='store_true')
    parser.add_argument(CmdArgs.num_trials_optional, help=CmdArgs.num_trials_optional_help, type=int, default=1)
    parser.add_argument(CmdArgs.for_all_matrices_optional, help=CmdArgs.for_all_matrices_optional_help, action = 'store_true')

    args = parser.parse_args(args = None if sys.argv[1:] else ['--help'])
    args_dict = vars(args)

    execute_from_command_line(
        args_dict.get(CmdArgs.matrix_size[2:]),
        wc_and_nwc = args_dict.get(CmdArgs.wc_and_nwc_optional.strip()[2:]),
        verbose = args_dict.get(CmdArgs.verbose_optional.strip()[2:]),
        max_iterations = args_dict.get(CmdArgs.max_iterations_optional.strip()[2:]),
        num_threads = args_dict.get(CmdArgs.num_threads_optional.strip()[2:]),
        on_plm = args_dict.get(CmdArgs.on_plm_optional.strip()[2:]),
        num_trials = args_dict.get(CmdArgs.num_trials_optional.strip()[2:]),
        for_all_matrices = args_dict.get(CmdArgs.for_all_matrices_optional.strip()[2:]),
    )

if __name__ =='__main__':
    train_coconet()


"""
Representation of the quantites, that can be added by the WSMs:

retrievalAddAbsSpecies
retrievalAddFreqShift
retrievalAddFreqStretch
retrievalAddConstantVMRAbsSpecies
retrievalAddCatalogParameter
retrievalAddCatalogParameters
retrievalAddMagField
retrievalAddPointingZa
retrievalAddPolyfit
retrievalAddScatSpecies
retrievalAddSinefit
retrievalAddSpecialSpecies
retrievalAddWind
retrievalAddSurfaceQuantity
retrievalAddTemperature
"""

import numpy as np
from scipy import sparse
import xarray as xr
import pandas as pd
from contextlib import contextmanager

from .boilerplate import set_variable_by_xml

from retrievals.arts import boilerplate
from retrievals.arts.atmosphere import p2z_simple, z2p_simple
from retrievals import level2

import warnings

@contextmanager
def retrieval_def(ws):
    """Context manager for the RetrievalDef[Init/Close] context."""
    ws.covmat_block = []
    ws.covmat_inv_block = []

    ws.retrievalDefInit()
    yield ws
    ws.retrievalDefClose()

def _covmat_shape(g1, g2, g3):
    n = len(g1) * len(g2) * len(g3)
    return (n, n)

def _jqs_sizes(jqs):
    """Get the total size of the jacobian quantities."""
    jq_size = list()
    for jq in jqs:
        sz = 1
        for g in jq.grids:
            sz *= len(g)
        jq_size.append(sz)
    return jq_size

def _sizes_to_slices(jq_size):
    """Convert the sizes to start and end indices."""
    jq_start = [0] + list(np.cumsum(jq_size)[:-1])
    jq_slices = [slice(start, start + size) for start, size in zip(jq_start, jq_size)]
    return jq_slices

def _jqs_slices(jqs):
    """Get the start and end indices of jacobian quantities."""
    return _sizes_to_slices(_jqs_sizes(jqs))

class RetrievalQuantity:
    """A generic retrieval quantity."""

    def __init__(self, kind, covmat, **kwargs):
        """
        :param kind: The kind is used to identify the corresponding WSMs, for example for a kind of `Sinefit` the WSM
                     `retrievalAddSinefit` would be called upon apply with the keyword arguments.
        :param covmat: The covariance matrix for the retrieval.
        :param kwargs: Arguments to the `retrievalAdd...` WSM.
        """
        self._kind = kind
        self._covmat = covmat
        self._args = kwargs

        self._is_applied = False
        self._sequence_id = None  # index inside jacobian quantities of workspace
        self._jacobian_quantity = None
        self._slice = None  # Holds the start and end indices for this quantity
        self._ws = None  # Associated workspace

        self._xa = None
        self._x = None
        self._avkm = None
        self._eo = None
        self._es = None

    def apply(self, ws):
        """
        Add this retrieval quantity to a Workspace and extract the indices from the jacobian quantities.
        """
        if self._is_applied:
            raise Exception('Same retrieval quantity cannot be added twice to the same workspace.')

        self._ws = ws
        self.apply_retrieval()
        self.apply_covmat()

        jqs = ws.jacobian_quantities.value
        self._sequence_id = len(jqs) - 1
        self._jacobian_quantity = jqs[self._sequence_id]
        self._slice = _jqs_slices(jqs)[self._sequence_id]
        self._is_applied = True

    def apply_retrieval(self):
        """Calls the corresponding `retrievalAdd...()` WSM of the workspace."""
        ws = self.ws
        boilerplate.clear_sparse(ws, ws.covmat_block)
        boilerplate.clear_sparse(ws, ws.covmat_inv_block)
        retrieval_add_wsm = getattr(ws, 'retrievalAdd' + self.kind)
        retrieval_add_wsm(**self.args)

    def apply_covmat(self):
        """Add the corresponding covmat block to the Workspace."""
        ws = self.ws
        set_variable_by_xml(ws, ws.covmat_block, self.covmat)
        ws.covmat_sxAddBlock(block=ws.covmat_block)

    def extract_apriori(self, xa=None):
        """
        Extract the a priori values. Call WSM `xaStandard()` before calling this function!
        :param xa: The values in the `xa` vector just before the inversion. If None, copy from Workspace.
        """
        ws = self._ws
        if xa is None:
            xa = ws.xa.value
        self._xa = xa[self._slice]

    def extract_result(self, x=None, avk=None, eo=None, es=None):
        """
        Extract the result and the corresponding values form the Workspace.
        If the arguments are `None`, the variables are copied from the associated Workspace.
        """
        ws = self._ws
        if x is None:
            x = ws.x.value
        if avk is None:
            avk = ws.avk.value
        if eo is None:
            eo = ws.retrieval_eo.value
        if es is None:
            es = ws.retrieval_ss.value

        self._x = x[self._slice]
        self._avkm = avk[self._slice, self._slice]
        self._eo = eo[self._slice]
        self._es = es[self._slice]

    def to_xarray(self):
        """
        Export this retrieval quantity.

        :rtype: xarray.Dataset
        """
        if '-' in self.slug or ' ' in self.slug or ',' in self.slug:
            prefix = self.slug.replace(' ','_')
            prefix = prefix.replace(',','_')
            prefix = prefix.replace('-','_')
            prefix = prefix.replace('__','_')
            prefix = prefix + '_'
        else:
            prefix = self.slug + '_'
        shape = self.shape
        grid_names = [prefix + gn for gn in self.grid_names]
        grids = self.grids
        flat_grid_name = prefix + 'grid' if self.dimensions > 1 else grid_names[0]

        coords = {n: c for n, c in zip(grid_names, grids)}

        ds = xr.Dataset(
            data_vars={
                prefix + 'x': (grid_names, np.reshape(self.x, shape, order='F')),
                prefix + 'xa': (grid_names, np.reshape(self.xa, shape, order='F')),
                prefix + 'mr': (grid_names, np.reshape(self.mr, shape, order='F')),
                prefix + 'eo': (grid_names, np.reshape(self.eo, shape, order='F')),
                prefix + 'es': (grid_names, np.reshape(self.es, shape, order='F')),
                prefix + 'avkm': ((flat_grid_name, flat_grid_name + '_avk'), self.avkm),
            },
            coords=coords,
            attrs={
                'maintag': self._jacobian_quantity.maintag,
                'subtag': self._jacobian_quantity.subtag,
                'subsubtag': self._jacobian_quantity.subsubtag,
                'analytical': self._jacobian_quantity.analytical,
                'mode': self._jacobian_quantity.mode,
                'perturbation': self._jacobian_quantity.perturbation,
            }
        )
        return ds

    @property
    def kind(self):
        return self._kind

    @property
    def covmat(self):
        return self._covmat

    @property
    def args(self):
        """The arguments for the `retrievalAdd...` WSM."""
        return self._args

    @property
    def num_elem(self):
        """Number of elements, product of all grid lengths."""
        return self.covmat.shape[0]

    @property
    def grids(self):
        """Associated grids."""
        return self._jacobian_quantity.grids

    @property
    def grid_names(self):
        """Names of the grids."""
        return ['grid' + str(i + 1) for i in range(len(self.shape))]

    @property
    def shape(self):
        """Shape of grids."""
        return tuple(map(len, self.grids))

    @property
    def dimensions(self):
        return len(list(filter(lambda x: x > 1, self.shape)))

    @property
    def ws(self):
        """Associated workspace. Set by :py:meth:`apply`."""
        return self._ws

    def __str__(self):
        return self.kind

    @property
    def slug(self):
        """A slug according to the kind of this retrieval quantity, useful for serialisation."""
        jq = self._jacobian_quantity
        if jq.maintag == 'Absorption species':
            return str.lower(jq.subtag)  # this is the species e.g. O3
        elif jq.maintag == 'Wind':
            return 'wind_' + str.lower(jq.subtag)  # component u, v, w
        elif jq.maintag == 'Polynomial baseline fit':
            return 'polyfit'  # Same for all coefficients
        else:
            return str.lower('_'.join((jq.maintag, jq.subtag, jq.subsubtag))).replace(' ', '-')

    @property
    def x(self):
        """The retrieved values. Available after :py:meth:`extract_results`."""
        return self._x

    @property
    def xa(self):
        """The a priori values. Available after :py:meth:`extract_apriori`."""
        return self._xa

    @property
    def avkm(self):
        """The averaging kernel matrix. Available after :py:meth:`extract_results`."""
        return self._avkm

    @property
    def mr(self):
        """The measurement response. Available after :py:meth:`extract_results`."""
        return level2.avkm_mr(self._avkm)

    @property
    def eo(self):
        """The retrieval error associated with the observational system. Available after :py:meth:`extract_results`."""
        return self._eo

    @property
    def es(self):
        """The smoothing error. Available after :py:meth:`extract_results`."""
        return self._es

class GriddedRetrievalQuantity(RetrievalQuantity):
    """Retrieval quantity that is retrieved on a spatial grid."""

    def __init__(self, kind, p_grid, lat_grid, lon_grid, covmat, **kwargs):
        """
        :param kind: The kind is used to identify the corresponding WSMs, for example for a kind of `Sinefit` the WSM
                     `retrievalAddSinefit` would be called upon apply with the keyword arguments.
        :param p_grid: Pressure grid
        :param lat_grid: Latitude grid
        :param lon_grid: Longitude grid
        :param covmat: The covariance matrix for the retrieval.
        :param kwargs: Arguments to the `retrievalAdd...` WSM.
        """
        kwargs['g1'] = p_grid
        kwargs['g2'] = lat_grid
        kwargs['g3'] = lon_grid
        if not covmat.shape == _covmat_shape(p_grid, lat_grid, lon_grid):
            expected = _covmat_shape(p_grid, lat_grid, lon_grid)
            raise ValueError(
                f'Covariance matrix must have shape according to retrieval grid elements: {expected[0]} x {expected[1]}')
        super().__init__(kind, covmat, **kwargs)

    def to_xarray(self):
        prefix = self.slug + '_'
        shape = self.shape
        grid_names = [prefix + gn for gn in self.grid_names]

        ds = super().to_xarray()

        if self.dimensions == 1:
            ds[prefix + 'z'] = (grid_names[0], self.z_grid)
            ds[prefix + 'fwhm'] = (grid_names, np.reshape(self.fwhm, shape, order='F'))
            ds[prefix + 'offset'] = (grid_names, np.reshape(self.offset, shape, order='F'))

        return ds

    @property
    def shape(self):
        return tuple(map(len, (self.args['g1'], self.args['g2'], self.args['g3'])))

    @property
    def num_elem(self):
        shape = self.shape
        return shape[0] * shape[1] * shape[2]

    @property
    def p_grid(self):
        return self._args['g1']

    @property
    def z_grid(self):
        if self.dimensions > 1:
            warnings.warn('Z grid extraction not supported for multi dim retrieval grids. Using simple conversion.')
            return p2z_simple(self.p_grid)

        atmosphere_dim = self._ws.atmosphere_dim.value
        lat_grid = self._ws.lat_grid.value
        lon_grid = self._ws.lon_grid.value
        if atmosphere_dim > 1:
            if self.lat_grid[0] not in lat_grid or self.lon_grid[0] not in lon_grid:
                warnings.warn('Z grid extraction not supported for arbitrary retrieval grids. Using simple conversion.')
                return p2z_simple(self.p_grid)

        z_field = self._ws.z_field.value

        if atmosphere_dim == 1:
            i_lat = i_lon = 0
        elif atmosphere_dim == 2:
            i_lon = 0
            i_lat = np.searchsorted(lat_grid, self.lat_grid[0])
        else:
            i_lat = np.searchsorted(lat_grid, self.lat_grid[0])
            i_lon = np.searchsorted(lon_grid, self.lon_grid[0])
        z_grid1 = z_field[:, i_lat, i_lon]
        p_grid1 = self._ws.p_grid.value
        idx = np.argsort(p_grid1)
        z_grid = np.interp(np.log(self.p_grid), np.log(p_grid1[idx]), z_grid1[idx])
        return z_grid

    @property
    def lat_grid(self):
        return self._args['g2']

    @property
    def lon_grid(self):
        return self._args['g3']

    @property
    def grid_names(self):
        return 'p', 'lat', 'lon'

    @property
    def fwhm(self):
        if self.dimensions > 1:
            raise NotImplementedError('FWHM not implemented for multi dim retrieval grids.')
        return level2.avkm_fwhm(self.z_grid, self.avkm)

    @property
    def offset(self):
        if self.dimensions > 1:
            raise NotImplementedError('Offset not implemented for multi dim retrieval grids.')
        return level2.avkm_offset(self.z_grid, self.avkm)

class AbsSpecies(GriddedRetrievalQuantity):
    """Absorption species."""

    # def __init__(self, species, p_grid, lat_grid, lon_grid, covmat,
    #              method='analytical', unit='rel', for_species_tag=1, dx=0.001):
    def __init__(self, species, p_grid, lat_grid, lon_grid, covmat,
                 unit='rel', for_species_tag=1):
        super().__init__('AbsSpecies',
                         p_grid=p_grid,
                         lat_grid=lat_grid,
                         lon_grid=lon_grid,
                         covmat=covmat,
                         species=species,
                         #method=method,
                         unit=unit,
                         for_species_tag=for_species_tag)#,
                         #dx=dx)

    @property
    def species(self):
        return self.args['species']

    @property
    def slug(self):
        return self.species.lower()

class Wind(GriddedRetrievalQuantity):
    """Wind."""

    def __init__(self, component, p_grid, lat_grid, lon_grid, covmat, dfrequency=0.1):
        super().__init__('Wind',
                         p_grid=p_grid,
                         lat_grid=lat_grid,
                         lon_grid=lon_grid,
                         covmat=covmat,
                         component=component,
                         dfrequency=dfrequency)

    @property
    def component(self):
        return self.args['component']

    @property
    def slug(self):
        return 'wind_' + self.component.lower()

class FreqShift(RetrievalQuantity):
    """Backend frequency shift."""

    def __init__(self, var, df=100e3):
        covmat = np.array([var ** 2])
        super().__init__('FreqShift', covmat, df=df)

    def apply_covmat(self):
        ws = self.ws
        ws.covmat_block = sparse.bsr_matrix(self.covmat)
        ws.covmat_sxAddBlock(block=ws.covmat_block)

    @property
    def num_elem(self):
        return 1

    @property
    def slug(self):
        return 'freq_shift'

class Polyfit(RetrievalQuantity):
    """Polynomial baseline fit."""

    def __init__(self, poly_order, covmats, pol_variation=True, los_variation=True, mblock_variation=True):
        if len(covmats) != poly_order + 1:
            raise ValueError('Must provide (poly_order + 1) covariance matrices.')
        for covmat in covmats:
            if covmat.ndim != 2:
                raise ValueError('Covariance matrices must have ndim=2, but got {}.'.format(covmat.ndim))

        self._covmats = covmats
        super().__init__('Polyfit', None,
                         poly_order=poly_order,
                         no_pol_variation=0 if pol_variation else 1,
                         no_los_variation=0 if los_variation else 1,
                         no_mblock_variation=0 if mblock_variation else 1)

    def apply(self, ws):
        if self._is_applied:
            raise Exception('Same retrieval quantity cannot be added twice to the same workspace.')

        self._ws = ws
        self.apply_retrieval()
        self.apply_covmat()

        n_jqs = self.poly_order + 1
        jqs = ws.jacobian_quantities.value
        self._sequence_id = slice(len(jqs) - n_jqs, len(jqs))
        self._jacobian_quantity = jqs[self._sequence_id]
        self._slice = _jqs_slices(jqs)[self._sequence_id]
        self._is_applied = True

    def apply_retrieval(self):
        ws = self.ws
        boilerplate.clear_sparse(ws, ws.covmat_block)
        boilerplate.clear_sparse(ws, ws.covmat_inv_block)
        ws.retrievalAddPolyfit(**self.args)

    def apply_covmat(self):
        for covmat in self._covmats:
            self.ws.covmat_sxAddBlock(block=covmat)

    def extract_apriori(self, xa=None):
        """
        Extract the a priori values. Call WSM `xaStandard()` before calling this function!
        :param xa: The values in the `xa` vector just before the inversion. If None, copy from Workspace.
        """
        ws = self._ws
        if xa is None:
            xa = ws.xa.value
        self._xa = [xa[s] for s in self._slice]  # Always 0, but we implement it anyways.

    def extract_result(self, x=None, avk=None, eo=None, es=None):
        """Extract the result and the corresponding values form the Workspace."""
        ws = self._ws
        if x is None:
            x = ws.x.value
        if avk is None:
            avk = ws.avk.value
        if eo is None:
            eo = ws.retrieval_eo.value
        if es is None:
            es = ws.retrieval_ss.value

        self._x = [x[s] for s in self._slice]
        self._avkm = [avk[s, s] for s in self._slice]
        self._eo = [eo[s] for s in self._slice]
        self._es = [es[s] for s in self._slice]

    def to_xarray(self):
        prefix = self.slug + '_'
        grid_names = ('poly_order', 'observation')
        grids = self.grids
        flat_grid_name = prefix + 'grid'

        coords = {n: c for n, c in zip(grid_names, grids)}

        ds = xr.Dataset(
            data_vars={
                prefix + 'x': (grid_names, np.stack(self.x)),
                prefix + 'xa': (grid_names, np.stack(self.xa)),
                prefix + 'mr': (grid_names, np.stack(self.mr)),
                prefix + 'eo': (grid_names, np.stack(self.eo)),
                prefix + 'es': (grid_names, np.stack(self.es)),
                prefix + 'avkm': ((grid_names[0], flat_grid_name, flat_grid_name + '_avk'), np.stack(self.avkm)),
            },
            coords=coords,
            attrs={
                'maintag': self._jacobian_quantity[0].maintag,
                'subtag': '',
                'subsubtag': '',
                'analytical': self._jacobian_quantity[0].analytical,
                'mode': self._jacobian_quantity[0].mode,
                'perturbation': self._jacobian_quantity[0].perturbation,
            }
        )
        return ds

    @property
    def slug(self):
        return 'poly_fit'

    @property
    def grids(self):
        jq_grids = self._jacobian_quantity[0].grids
        if len(jq_grids[0]) > 1 or len(jq_grids[1]) > 1 or len(jq_grids[2]) > 1:
            raise NotImplementedError()
        return [np.arange(self.poly_order + 1), np.array(jq_grids[-1], dtype=np.int)]

    @property
    def covmat(self):
        return sparse.diags(self._covmats)

    @property
    def poly_order(self):
        return self.args['poly_order']

    @property
    def mr(self):
        return [level2.avkm_mr(a) for a in self._avkm]

class Sinefit(RetrievalQuantity):
    """Sinusoidal baseline fit."""

    def __init__(self, periods, covmats, pol_variation=True, los_variation=True, mblock_variation=True):
        # if len(covmats) != poly_order + 1:
        #     raise ValueError('Must provide (poly_order + 1) covariance matrices.')
        # for covmat in covmats:
        #     if covmat.ndim != 2:
        #         raise ValueError('Covariance matrices must have ndim=2, but got {}.'.format(covmat.ndim))

        self._covmats = covmats
        super().__init__('Sinefit', None,
                         periods=periods,
                         no_pol_variation=0 if pol_variation else 1,
                         no_los_variation=0 if los_variation else 1,
                         no_mblock_variation=0 if mblock_variation else 1)

    def apply(self, ws):
        if self._is_applied:
            raise Exception('Same retrieval quantity cannot be added twice to the same workspace.')

        self._ws = ws
        self.apply_retrieval()
        self.apply_covmat()

        n_jqs = self.poly_order + 1
        jqs = ws.jacobian_quantities.value
        self._sequence_id = slice(len(jqs) - n_jqs, len(jqs))
        self._jacobian_quantity = jqs[self._sequence_id]
        self._slice = _jqs_slices(jqs)[self._sequence_id]
        self._is_applied = True

    def apply_retrieval(self):
        ws = self.ws
        boilerplate.clear_sparse(ws, ws.covmat_block)
        boilerplate.clear_sparse(ws, ws.covmat_inv_block)
        ws.retrievalAddSinefit(**self.args)

    def apply_covmat(self):
        for covmat in self._covmats:
            self.ws.covmat_sxAddBlock(block=covmat)

    def extract_apriori(self, xa=None):
        """
        Extract the a priori values. Call WSM `xaStandard()` before calling this function!
        :param xa: The values in the `xa` vector just before the inversion. If None, copy from Workspace.
        """
        ws = self._ws
        if xa is None:
            xa = ws.xa.value
        self._xa = [xa[s] for s in self._slice]  # Always 0, but we implement it anyways.

    def extract_result(self, x=None, avk=None, eo=None, es=None):
        """Extract the result and the corresponding values form the Workspace."""
        ws = self._ws
        if x is None:
            x = ws.x.value
        if avk is None:
            avk = ws.avk.value
        if eo is None:
            eo = ws.retrieval_eo.value
        if es is None:
            es = ws.retrieval_ss.value

        self._x = [x[s] for s in self._slice]
        self._avkm = [avk[s, s] for s in self._slice]
        self._eo = [eo[s] for s in self._slice]
        self._es = [es[s] for s in self._slice]

    def to_xarray(self):
        prefix = self.slug + '_'
        grid_names = ('poly_order', 'observation')
        grids = self.grids
        flat_grid_name = prefix + 'grid'

        coords = {n: c for n, c in zip(grid_names, grids)}

        ds = xr.Dataset(
            data_vars={
                prefix + 'x': (grid_names, np.stack(self.x)),
                prefix + 'xa': (grid_names, np.stack(self.xa)),
                prefix + 'mr': (grid_names, np.stack(self.mr)),
                prefix + 'eo': (grid_names, np.stack(self.eo)),
                prefix + 'es': (grid_names, np.stack(self.es)),
                prefix + 'avkm': ((grid_names[0], flat_grid_name, flat_grid_name + '_avk'), np.stack(self.avkm)),
            },
            coords=coords,
            attrs={
                'maintag': self._jacobian_quantity[0].maintag,
                'subtag': '',
                'subsubtag': '',
                'analytical': self._jacobian_quantity[0].analytical,
                'mode': self._jacobian_quantity[0].mode,
                'perturbation': self._jacobian_quantity[0].perturbation,
            }
        )
        return ds

    @property
    def slug(self):
        return 'poly_fit'

    @property
    def grids(self):
        jq_grids = self._jacobian_quantity[0].grids
        if len(jq_grids[0]) > 1 or len(jq_grids[1]) > 1 or len(jq_grids[2]) > 1:
            raise NotImplementedError()
        return [np.arange(self.poly_order + 1), np.array(jq_grids[-1], dtype=np.int)]

    @property
    def covmat(self):
        return sparse.diags(self._covmats)

    @property
    def poly_order(self):
        return self.args['poly_order']

    @property
    def mr(self):
        return [level2.avkm_mr(a) for a in self._avkm]

class SlicedArray:
    def __init__(self, values, slugs_to_slices):
        self._values = values
        self.slugs = slugs_to_slices

    def __getitem__(self, item):
        if isinstance(item, tuple):
            if len(item) == len(self._values.shape):
                idx = tuple(slice(*self.slugs[it]) for it in item)
                return self.values[idx]
            else:
                raise ValueError('Number of indexers must be one or be equal to number of dimensions.')
        else:
            a, b = self.slugs[item]
            n_dims = len(self._values.shape)
            return self._values[tuple(n_dims * [slice(a, b)])]

    @property
    def values(self):
        return self._values