#!/usr/bin/env python
# encoding: utf-8

"""
File: cloudforest/core.py

Created by Nicholas G. Crawford and Brant C. Faircloth
Copyright (c) 2011-2012 Nicholas G. Crawford and Brant C. Faircloth. All rights reserved.

Description: Core functions for cloudforest

"""

import os
import re
import sys
import shutil
import argparse
import platform
import tempfile
import subprocess
import numpy as np

#import pdb


def is_dir(dirname):
    if not os.path.isdir(dirname):
        msg = "{0} is not a directory".format(dirname)
        raise argparse.ArgumentTypeError(msg)
    else:
        return dirname


def is_file(filename):
    if not os.path.isfile:
        msg = "{0} is not a file".format(filename)
        raise argparse.ArgumentTypeError(msg)
    else:
        return filename


def format_oneliner_from_dict(taxa_seq_dict, name, model=None):
    assert len(taxa_seq_dict) > 0, ValueError("Phylip file has no content")
    seq = ['%s,%s' % (k, v) for k, v in taxa_seq_dict.iteritems()]
    if not model:
        return "chrm=%s:%s;" % (name.strip(), ','.join(seq))
    else:
        return "chrm=%s,model=%s:%s;" % (name.strip(), model, ','.join(seq))


def phylip_to_oneliner(phylip, locus, model=None):
    taxa_id_dict, taxa_seq_dict = {}, {}
    count = 0
    for lineno, line in enumerate(phylip.split('\n')):
        # skip blank lines
        if len(line.strip()) == 0:
            continue
        if lineno == 0:
            taxa_count, align_len = [int(i) for i in line.strip().split()]
        # INITIALIZE DICTS WITH TAXA ID'S AND INITIAL SEQS
        if 0 < count <= taxa_count:
            # clean up errant spaces
            sp = [val for val in line.split(' ') if val != '']
            name = sp[0]
            sequence = ''.join(sp[1:])
            taxa_id_dict[count] = name
            taxa_seq_dict[name] = sequence.replace(' ', '')
        # ADD ADDITIONAL LINES TO ALIGNMENT
        if count > taxa_count:
            num = count % taxa_count
            if num == 0:
                num = max(taxa_id_dict.keys())
            name = taxa_id_dict[num]
            taxa_seq_dict[name] += line.strip().replace(' ', '')
        count += 1
    oneliner = format_oneliner_from_dict(taxa_seq_dict, locus, model)
    return oneliner


def oneliner_to_phylip(line):
    """Convert one-liner to phylip format."""

    # Problem: Phylip has a 10 char limit that is not accounted for. (NGC)
    # "Strict Phylip expects the first character state to appear on Column 11
    #  for each and every sequence, no ifs, and, or buts."
    # http://www.phylo.org/tools/phylip.html

    # Remove additional info (e.g., chrm,model,etc.) from line (fixed above problem)
    if ':' in line:
        line = line.split(":")[-1]
    seqs = line.strip(";\n").split(',')
    label_seqs = zip(seqs[:-1:2], seqs[1::2])
    #print label_seqs
    taxa_count = len(label_seqs)
    seq_length = len(label_seqs[0][1])
    # pad all names to length of longest name + 1 space
    max_name_length = max([len(val) for val in seqs[:-1:2]]) + 1
    # add header
    header = "%s %s\n" % (taxa_count, seq_length)
    alignment = '\n'.join(['%s%s' % (i[0].ljust(max_name_length), i[1]) for i in label_seqs])
    return header + alignment


def get_bootstraps(sample, replicates=1, return_choices=False):
    """Create boostrapped replicates of an a numpy array.

    Returns list

    """
    # make sure we instantiate a Random number generator that doesn't
    # fall victim to the forking issues w/ multiprocessing
    #
    # Create generator w/ seed.  None as input draws from /dev/urandom
    rng = np.random.RandomState()
    if not isinstance(sample, np.ndarray):
        try:
            sample = np.array(sample)
        except:
            raise TypeError("bootstrap() input must be list or numpy.ndarray")
    size = len(sample)
    if replicates == 1:
        choices = rng.randint(0, size, size)
        # for debugging
        #sys.stdout.write(str(choices))
        if not return_choices:
            return sample[choices].tolist()
        else:
            return sample[choices].tolist(), choices

    else:
        return [sample[rng.randint(0, size, size)].tolist()
                        for i in xrange(replicates)]


def oneliner_to_array(line):
    """Convert oneliner to 2d numpy array.

    Returns tuple(list, array)

    """
    seqs = line.split(",")
    label_seqs = zip(seqs[:-1:2], seqs[1::2])
    taxa, bases = [], []
    for taxon, seq in label_seqs:
        bases.append(list(seq))
        taxa.append(taxon.strip())
    return taxa, np.array(bases)


def array_to_oneliner(taxa, bases):
    """Convert array of array of taxa and an array of bases to one-liner.

    Returns string

    """
    # join is faster and cleaner than concatenate
    # list comp is faster than for
    # enumerate gets index of taxon from taxa
    oneliner = ','.join([','.join([taxa[count], seq.tostring()]) for count, seq in enumerate(bases)])
    return oneliner


def make_tree_name(args_dict):
    """Converts dictionary of arguement into a sorted string with the
    format:  """
    # join is faster and cleaner than concatenate
    # list comp is faster than for
    name = ','.join(['='.join([pair[0], pair[1]]) for pair in sorted(args_dict.items())])
    return name


def split_oneliner(line, args_dict=None, default_model=False):
    """From a oneliner, split locus/taxon info into dict and locus into string

    Returns dict or tuple(dict, locus)

    """
    if not args_dict:
        args_dict = {}
    if ':' in line:
        args, locus = line.split(":")
        args = args.split(',')
        for item in args:
            dict_key, value = item.split("=")
            args_dict[dict_key] = value
    else:
        locus = line
    if default_model:
        if 'model' not in args_dict.keys():
            args_dict['model'] = 'GTR'
    return args_dict, locus


class FullPaths(argparse.Action):
    """Expand user- and relative-paths"""
    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))


class Process():
    """ """
    def __init__(self):
        pass

    def get_bootstrap_replicates(self, key, line):
        loci = line.strip().split(';')
        loci = loci[:-1]
        # first, bootstrap across loci
        multilocus_bstrap = get_bootstraps(loci)
        # second, bootstrap bases within loci
        for locus in multilocus_bstrap:
            # split keys from alignments
            args_dict, locus = split_oneliner(locus)
            # convert alignments to arrays
            taxa, numpy_alignment = oneliner_to_array(locus)
            # transpose so we bootstrap by columns
            bases_by_col = np.column_stack(numpy_alignment)
            # bootstrap by columns
            shuffled = get_bootstraps(bases_by_col)
            # transpose back to rows of sequences
            shuffled = np.column_stack(shuffled)
            # convert back to oneliner
            oneliner = array_to_oneliner(taxa, shuffled)
            oneliner = "%s:%s" % (make_tree_name(args_dict), oneliner)
            yield key, oneliner

    def process_stats_file(self, fin):
        """"given an input phyml stats file, return the log-likelihood of the tree"""
        result = None
        regex = re.compile("Log-likelihood:\s+(.+)")
        for line in fin:
            result = regex.search(line)
            if result:
                break
        if result is None:
            raise ValueError("No Log-likelihood found")
        return float(result.groups()[0])

    def get_genetrees(self, key, line, pth='bin', genetrees=True, no_model=False):
        """Compute genetrees using model in oneliner. Parses out evolutionary