suppressMessages(library("argparser"))
suppressMessages(library("vegan"))


# Sets up an argparser to make this tractable
parser <- arg_parser(
  description="Basic Adonis Functionality for age and set"
)
parser <- add_argument(parser,
  '--distance-matrix', 
  help='the distance matrix',
)
parser <- add_argument(parser,
  '--metadata', 
  help=('Path to the metadata file'),
)
parser <- add_argument(parser,
  '--output',
  help=('The location for the summary file')
)

args <- parse_args(parser)

# reads in the distance matrix as a dataframe and sets the rownames as the first column
df <- read.csv2(args$distance_matrix, sep='\t', row.names='X')
map <- read.csv2(args$metadata, sep='\t', row.names='sample.id')
map <- map[row.names(df),]

dm  <- dist(df[row.names(df), row.names(df)])

age <- adonis2(dm ~ age, data=map)

res <- data.frame(
  age = c(age$R2[1], age$`Pr(>F)`[1], age$F[1]),
  row.names = c('R2', 'p_999', 'F')
  )

write.table(res, file=args$output, quote=FALSE, sep='\t', col.names = NA)

description = """
This script will build a set of synthetic samples based on resampling of the
american gut data. We'll do this by reading in the biom table and then 
randomly selecting a subset of samples using non ambigious samples. 

The samples will then be split into two approximately equal sized groups and 
used to generate a set of normalized, simulated samples
"""

import argparse
import os

import biom
import numpy as np
import pandas as pd

from qiime2 import Artifact


def subsample_data(metadata, table, age_split_fun, id_prefixes,
                   colors=None,
                   group_size=10, 
                   n_samples=20, seed=1234):
    """
    Generates a set of synthetic sequences based on the American Gut using
    seperation between people in their 20s and their 60s

    Parameters
    ----------
    keep_map : DataFrame
        A map of the samples with the age encoded in `age_cat` and the 
        index as the sample number
    samples : biom.Table
        The samples used for simulation
    age_split_fun: function
    group_size : int
        The number of samples to be combined to make the simulated sample
    split_size : int
        The number of samples to use in each set. (This is used to adjust)
        for uneven group sizes in the two age groups
    n_samples : int
        The number of samples to simulate for each age group

    Return 
    ------
    biom.Table
        All of the synthesized samples

    """
    # Sets a seed for consistency
    np.random.seed(seed)

    # Gets the age catgories 
    metadata['age_cat'] = metadata['age'].apply(age_split_fun)
    metadata.dropna(subset=['age_cat'], inplace=True)

    table.filter(metadata.index, axis='sample', inplace=True)
    table.add_metadata(metadata[['age_cat']].to_dict(orient='index'), axis='sample')

    print(id_prefixes)

    tables = [
        table.copy().filter(lambda v, id_, md: md['age_cat'] == prefix,
                            axis='sample', inplace=False)
        for prefix in id_prefixes
    ]

    synthetic = [
        tidy_names(pd.concat(axis=1, objs=[
            build_subsample(table, group_size, prefix) 
            for i in range(n_samples)]), 
        prefix)
        for prefix, table in zip(*(id_prefixes, tables))
    ]
    print(synthetic[0].head())

    all_samples = pd.concat(axis=1, sort=False, objs=synthetic)
    all_samples.fillna(0, inplace=True)

    all_biom = biom.Table(data=all_samples.values, 
                          sample_ids=all_samples.columns,
                          observation_ids=all_samples.index
                          )
    metadata = pd.DataFrame.from_dict(orient='index', data={
        id_: {'set': id_.split('.')[1],
              'age': id_.split('.')[2],
              }
        for id_ in all_biom.ids(axis='sample')
        })
    metadata.index.set_names('sample-id', inplace=True)
    metadata['composite'] = metadata.apply(lambda x: "{age}-{set}".format(**x),
                                            axis=1)
    if colors is not None:
        metadata['color'] = metadata['composite'].replace(colors)

    return all_biom, metadata

def tidy_names(table, prefix):
    table.rename(inplace=True, columns={
        int(i): f"sample.1.{prefix}.{i}" for i in table.columns
        })
    return table

def build_subsample(table, sample_size, group):
    """
    Builds a subsampled table from a biom object
    """
    sub_t = table.copy().subsample(sample_size, axis='sample', by_id=True)
    sub_t.filter(lambda v, id_, md: (v > 0).sum() > np.round(0.4 * sample_size), 
                 axis='observation', 
                 inplace=True)
    single = sub_t.copy().collapse(lambda id_, md: md['age_cat'], 
                                   axis='sample')
    single.norm(axis='sample', inplace=True)
    single.filter(lambda v, id_, md: (v > 2e-5).all(), 
                  axis='observation', 
                  inplace=True)
    single.norm(axis='sample', inplace=True)
    return pd.Series(np.round(single.data(group) * 5e5).round(0),
                     index=single.ids(axis='observation'))


parser = argparse.ArgumentParser(description=description)
parser.add_argument(
    '--keep-map',
    help=("The map of the samples to be used for simulation. This must "
          "contain  a sample identifier labeled `sample_name` and the age"
          "labeled age_cat"
          ),
    )
parser.add_argument(
    '--sample-table',
    help=("A biom table containing sample counts. The sample list should "
          "match those present in `--keep-map`"),
    )
parser.add_argument(
    '--simulated-table',
    help=('Filepath to the qiime artifact containing all the samples'),
    )
parser.add_argument(
    '--simulated-metadata',
    help=('Filepath to the qiime artifact containing all the samples'),
    )
parser.add_argument(
    '--samples-into-simulation',
    help=("The number of samples to combine into each simulated sample"),
    default=10,
    type=int,
    )
parser.add_argument(
    '--group-size',
    help=("The number of samples in each group to simulate for the set."),
    default=20,
    type=int,
    )
parser.add_argument(
    '--seed',
    help=('a random seed'),
    default=1234,
    type=int,
    )

if __name__ == '__main__':
    args = parser.parse_args()

    # Loads group
    keep_map = pd.read_csv(args.keep_map, sep='\t', dtype=str)   
    keep_map.set_index('sample_name', inplace=True)
    keep_map['age'] = keep_map['age'].astype(float)

    # Loads table
    sample_table = Artifact.load(args.sample_table).view(biom.Table)

    def cat_age(x):
        if pd.isnull(x):
            return x
        if x <= 3:
            return 'infant'
        elif (3 < x) & (x <= 12):
            return 'child'
        elif (20 < x) & (x <= 60):
            return 'adult'
        else:
            return np.nan

    colors = {'infant-1': '#1f78b4', 'infant-2': '#a6cee3', 
              'adult-1': '#e31a1c', 'adult-2': '#fb9a99'}


    all_biom, metadata = \
        subsample_data(metadata=keep_map.copy(), 
                       table=sample_table, 
                       age_split_fun=cat_age, 
                       id_prefixes=['infant', 'adult'],
                       group_size=args.samples_into_simulation,
                       n_samples=args.group_size,
                       seed=args.seed,
                       )
    all_table = Artifact.import_data('FeatureTable[Frequency]', all_biom)
    all_table.save(args.simulated_table)

    metadata.to_csv(args.simulated_metadata, sep='\t')

description = """
Simulates the an ASV table and dereplicated sequences for OTU clustering
and meta analysis exploration
"""


import argparse
import hashlib

import biom
import numpy as np
import pandas as pd
import regex
import skbio
from skbio import DNA

from qiime2 import Artifact


def build_regional_table(samples, seqs, read_length):
    """
    Amplifies a regional table from the original table

    Parameters
    ----------
    samples : DataFrame
        A biom-style dataframe where the index is sequences (corresponding)
        to the sequences in seqs

    """
    # Filters the sequences and table down to those present in the sequences
    degen_count = seqs.astype(str).apply(
        lambda x: len(regex.findall('[RYSWKMBDHVN]', str(x)))
        )
    len_check = seqs.astype(str).apply(lambda x: (len(x) >= np.absolute(read_length)))
    print(len_check.sum())
    degen_keep = degen_count.index[(degen_count <= 10) & len_check].values
    var_ids = [id_ for id_ in samples.index if (id_ in degen_keep)]
    samples = samples.loc[var_ids]

    table_seqs = seqs[var_ids].copy()

    # Expands the degenerate sequences
    exp_ = {id_: np.array([str(s) for s in seq_.expand_degenerates()])
            for id_, seq_ in table_seqs.iteritems()}

    # Generates the table by expanding the sequence where degenerates are
    # randomnly sampled and then combining the sequences into a table
    forward_reads = dict([])
    reverse_reads = dict([])
    joined_reads = dict([])
    for sample in samples.columns:
        expanded = pd.Series(np.hstack([
                np.random.choice(exp_[id_], size=count, replace=True)
                for id_, count in 
                samples.loc[samples[sample] > 0, sample].iteritems()])
        ).apply(lambda x: x[:read_length])
        joined_reads[sample] = expanded.value_counts()

    joined_table = pd.DataFrame.from_dict(orient='index', data=joined_reads).fillna(0)
    joined_table = joined_table.loc[joined_table.sum(axis=1) > 10]
    joined_reads = pd.Series({
        hash_seq(seq): DNA(seq) 
        for seq in joined_table.columns
    })
    joined_table.rename(columns={seq: hash_seq(seq) for seq in joined_table.columns},
                        inplace=True)
    jointable = Artifact.import_data('FeatureTable[Frequency]', joined_table, pd.DataFrame)
    joinseqs = Artifact.import_data('FeatureData[Sequence]', joined_reads, pd.Series)

    return jointable, joinseqs


def hash_seq(x):
    """
    A lazy function to hash things
    """
    return hashlib.md5(x.encode()).hexdigest()




parser = argparse.ArgumentParser(description=description)
parser.add_argument(
    '--sequences',
    help=('The extracted regional sequences'),
    )
parser.add_argument(
    '--sample-table',
    help=("A table of the samples to be simulated."),
    )
parser.add_argument(
    '--read-length',
    help=("length of the final simulated sequences. We treat these as paired"
          "end sequences"),
    type=int,
    )
parser.add_argument(
    '--table',
    help=('Path to save the joined-sequence table')
    )
parser.add_argument(
    '--reads',
    help=('Path to save the representative sequences for the joined table'),
    )
parser.add_argument(
    '--random',
    help=("A seed for randomn number generation"),
    default=None,
    )

if __name__ == '__main__':
    args = parser.parse_args()

    # Gets the regional sequences
    seqs =  Artifact.load(args.sequences).view(pd.Series)

    # Gets the table
    table = Artifact.load(args.sample_table).view(pd.DataFrame).fillna(0).astype(int).T

    # Simulates the samples
    if args.random is not None:
        np.random.seed(int(args.random))
    # fwd_table, rev_table, join_table, fwd_seqs, rev_seqs, join_seqs = \
    join_table, join_seqs = \
        build_regional_table(table, seqs, args.read_length)

    # fwd_table.save(args.fwd_table)
    # fwd_seqs.save(args.fwd_reads)
    # rev_table.save(args.rev_table)
    # rev_seqs.save(args.rev_reads)
    join_table.save(args.table)
    join_seqs.save(args.reads)

shell:
	"""
	qiime cutadapt trim-single \
	 --i-demultiplexed-sequences {input} \
	 --p-front {params.f_primer} \
	 --p-no-discard-untrimmed \
	 --o-trimmed-sequences {output}
	"""

shell:
	"""
	qiime vsearch join-pairs \
	--i-demultiplexed-seqs {input} \
	 --o-joined-sequences {output}
	"""