r"""

Util functions (:mod: `qiita_db.meta_util`)

===========================================

..currentmodule:: qiita_db.meta_util

This module provides utility functions that use the ORM objects. ORM objects

CANNOT import from this file.

Methods

-------

..autosummary::

    :toctree: generated/

    get_lat_longs

"""

# -----------------------------------------------------------------------------

# Copyright (c) 2014--, The Qiita Development Team.

#

# Distributed under the terms of the BSD 3-clause License.

#

# The full license is in the file LICENSE, distributed with this software.

# -----------------------------------------------------------------------------

from os import stat

from shutil import move

from os.path import join, relpath, basename

from time import strftime, localtime

import matplotlib.pyplot as plt

import matplotlib as mpl

from base64 import b64encode

from urllib.parse import quote

from io import BytesIO

from datetime import datetime

from collections import defaultdict, Counter

from tarfile import open as topen, TarInfo

from hashlib import md5

from re import sub

from json import loads, dump, dumps

from qiita_db.util import create_nested_path, retrieve_resource_data

from qiita_db.util import resource_allocation_plot

from qiita_core.qiita_settings import qiita_config, r_client

from qiita_core.configuration_manager import ConfigurationManager

import qiita_db as qdb

# global constant list used in resource_allocation_page

COLUMNS = [

    "sName", "sVersion", "cID", "cName", "processing_job_id",

    "parameters", "samples", "columns", "input_size", "extra_info",

    "MaxRSSRaw", "ElapsedRaw", "Start", "node_name", "node_model"]

RAW_DATA_ARTIFACT_TYPE = {

        'SFF', 'FASTQ', 'FASTA', 'FASTA_Sanger', 'per_sample_FASTQ'}

def _get_data_fpids(constructor, object_id):

    """Small function for getting filepath IDS associated with data object

    Parameters

    ----------

    constructor : a subclass of BaseData

        E.g., RawData, PreprocessedData, or ProcessedData

    object_id : int

        The ID of the data object

    Returns

    -------

    set of int

    """

    with qdb.sql_connection.TRN:

        obj = constructor(object_id)

        return {fpid for fpid, _, _ in obj.get_filepaths()}

def validate_filepath_access_by_user(user, filepath_id):

    """Validates if the user has access to the filepath_id

    Parameters

    ----------

    user : User object

        The user we are interested in

    filepath_id : int

        The filepath id

    Returns

    -------

    bool

        If the user has access or not to the filepath_id

    Notes

    -----

    Admins have access to all files so True is always returned

    """

    TRN = qdb.sql_connection.TRN

    with TRN:

        if user.level == "admin":

            # admins have access all files

            return True

        sql = """SELECT

            (SELECT array_agg(artifact_id)

             FROM qiita.artifact_filepath

             WHERE filepath_id = {0}) AS artifact,

            (SELECT array_agg(study_id)

             FROM qiita.sample_template_filepath

             WHERE filepath_id = {0}) AS sample_info,

            (SELECT array_agg(prep_template_id)

             FROM qiita.prep_template_filepath

             WHERE filepath_id = {0}) AS prep_info,

            (SELECT array_agg(analysis_id)

             FROM qiita.analysis_filepath

             WHERE filepath_id = {0}) AS analysis""".format(filepath_id)

        TRN.add(sql)

        arid, sid, pid, anid = TRN.execute_fetchflatten()

        # artifacts

        if arid:

            # [0] cause we should only have 1

            artifact = qdb.artifact.Artifact(arid[0])

            if artifact.visibility == 'public':

                # TODO: https://github.com/biocore/qiita/issues/1724

                if artifact.artifact_type in RAW_DATA_ARTIFACT_TYPE:

                    study = artifact.study

                    has_access = study.has_access(user, no_public=True)

                    if (not study.public_raw_download and not has_access):

                        return False

                return True

            else:

                study = artifact.study

                if study:

                    # let's take the visibility via the Study

                    return artifact.study.has_access(user)

                else:

                    analysis = artifact.analysis

                    return analysis in (

                        user.private_analyses | user.shared_analyses)

        # sample info files

        elif sid:

            # the visibility of the sample info file is given by the

            # study visibility

            # [0] cause we should only have 1

            return qdb.study.Study(sid[0]).has_access(user)

        # prep info files

        elif pid:

            # the prep access is given by it's artifacts, if the user has

            # access to any artifact, it should have access to the prep

            # [0] cause we should only have 1

            pt = qdb.metadata_template.prep_template.PrepTemplate(

                pid[0])

            a = pt.artifact

            # however, the prep info file could not have any artifacts attached

            # , in that case we will use the study access level

            if a is None:

                return qdb.study.Study(pt.study_id).has_access(user)

            else:

                if (a.visibility == 'public' or a.study.has_access(user)):

                    return True

                else:

                    for c in a.descendants.nodes():

                        if ((c.visibility == 'public' or

                             c.study.has_access(user))):

                            return True

            return False

        # analyses

        elif anid:

            # [0] cause we should only have 1

            aid = anid[0]

            analysis = qdb.analysis.Analysis(aid)

            return analysis.is_public | (analysis in (

                user.private_analyses | user.shared_analyses))

        return False

def update_redis_stats():

    """Generate the system stats and save them in redis

    Returns

    -------

    list of str

        artifact filepaths that are not present in the file system

    """

    STUDY = qdb.study.Study

    number_studies = {'public': 0, 'private': 0, 'sandbox': 0}

    number_of_samples = {'public': 0, 'private': 0, 'sandbox': 0}

    num_studies_ebi = 0

    num_samples_ebi = 0

    number_samples_ebi_prep = 0

    stats = []

    missing_files = []

    per_data_type_stats = Counter()

    for study in STUDY.iter():

        st = study.sample_template

        if st is None:

            continue

        # counting samples submitted to EBI-ENA

        len_samples_ebi = sum([esa is not None

                               for esa in st.ebi_sample_accessions.values()])

        if len_samples_ebi != 0:

            num_studies_ebi += 1

            num_samples_ebi += len_samples_ebi

        samples_status = defaultdict(set)

        for pt in study.prep_templates():

            pt_samples = list(pt.keys())

            pt_status = pt.status

            if pt_status == 'public':

                per_data_type_stats[pt.data_type()] += len(pt_samples)

            samples_status[pt_status].update(pt_samples)

            # counting experiments (samples in preps) submitted to EBI-ENA

            number_samples_ebi_prep += sum([

                esa is not None

                for esa in pt.ebi_experiment_accessions.values()])

        # counting studies

        if 'public' in samples_status:

            number_studies['public'] += 1

        elif 'private' in samples_status:

            number_studies['private'] += 1

        else:

            # note that this is a catch all for other status; at time of

            # writing there is status: awaiting_approval

            number_studies['sandbox'] += 1

        # counting samples; note that some of these lines could be merged with

        # the block above but I decided to split it in 2 for clarity

        if 'public' in samples_status:

            number_of_samples['public'] += len(samples_status['public'])

        if 'private' in samples_status:

            number_of_samples['private'] += len(samples_status['private'])

        if 'sandbox' in samples_status:

            number_of_samples['sandbox'] += len(samples_status['sandbox'])

        # processing filepaths

        for artifact in study.artifacts():

            for adata in artifact.filepaths:

                try:

                    s = stat(adata['fp'])

                except OSError:

                    missing_files.append(adata['fp'])

                else:

                    stats.append(

                        (adata['fp_type'], s.st_size, strftime('%Y-%m',

                         localtime(s.st_mtime))))

    num_users = qdb.util.get_count('qiita.qiita_user')

    num_processing_jobs = qdb.util.get_count('qiita.processing_job')

    lat_longs = dumps(get_lat_longs())

    summary = {}

    all_dates = []

    # these are some filetypes that are too small to plot alone so we'll merge

    # in other

    group_other = {'html_summary', 'tgz', 'directory', 'raw_fasta', 'log',

                   'raw_sff', 'raw_qual', 'qza', 'html_summary_dir',

                   'qza', 'plain_text', 'raw_barcodes'}

    for ft, size, ym in stats:

        if ft in group_other:

            ft = 'other'

        if ft not in summary:

            summary[ft] = {}

        if ym not in summary[ft]:

            summary[ft][ym] = 0

            all_dates.append(ym)

        summary[ft][ym] += size

    all_dates = sorted(set(all_dates))

    # sorting summaries

    ordered_summary = {}

    for dt in summary:

        new_list = []

        current_value = 0

        for ad in all_dates:

            if ad in summary[dt]:

                current_value += summary[dt][ad]

            new_list.append(current_value)

        ordered_summary[dt] = new_list

    plot_order = sorted([(k, ordered_summary[k][-1]) for k in ordered_summary],

                        key=lambda x: x[1])

    # helper function to generate y axis, modified from:

    # http://stackoverflow.com/a/1094933

    def sizeof_fmt(value, position):

        number = None

        for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:

            if abs(value) < 1024.0:

                number = "%3.1f%s" % (value, unit)

                break

            value /= 1024.0

        if number is None:

            number = "%.1f%s" % (value, 'Yi')

        return number

    all_dates_axis = range(len(all_dates))

    plt.locator_params(axis='y', nbins=10)

    plt.figure(figsize=(20, 10))

    for k, v in plot_order:

        plt.plot(all_dates_axis, ordered_summary[k], linewidth=2, label=k)

    plt.xticks(all_dates_axis, all_dates)

    plt.legend()

    plt.grid()

    ax = plt.gca()

    ax.yaxis.set_major_formatter(mpl.ticker.FuncFormatter(sizeof_fmt))

    plt.xticks(rotation=90)

    plt.xlabel('Date')

    plt.ylabel('Storage space per data type')

    plot = BytesIO()

    plt.savefig(plot, format='png')

    plot.seek(0)

    img = 'data:image/png;base64,' + quote(b64encode(plot.getbuffer()))

    time = datetime.now().strftime('%m-%d-%y %H:%M:%S')

    portal = qiita_config.portal

    # making sure per_data_type_stats has some data so hmset doesn't fail

    if per_data_type_stats == {}:

        per_data_type_stats['No data'] = 0

    vals = [

        ('number_studies', number_studies, r_client.hmset),

        ('number_of_samples', number_of_samples, r_client.hmset),

        ('per_data_type_stats', dict(per_data_type_stats), r_client.hmset),

        ('num_users', num_users, r_client.set),

        ('lat_longs', (lat_longs), r_client.set),

        ('num_studies_ebi', num_studies_ebi, r_client.set),

        ('num_samples_ebi', num_samples_ebi, r_client.set),

        ('number_samples_ebi_prep', number_samples_ebi_prep, r_client.set),

        ('img', img, r_client.set),

        ('time', time, r_client.set),

        ('num_processing_jobs', num_processing_jobs, r_client.set)]

    for k, v, f in vals:

        redis_key = '%s:stats:%s' % (portal, k)

        # important to "flush" variables to avoid errors

        r_client.delete(redis_key)

        f(redis_key, v)

    # preparing vals to insert into DB

    vals = dumps(dict([x[:-1] for x in vals]))

    sql = """INSERT INTO qiita.stats_daily (stats, stats_timestamp)

             VALUES (%s, NOW())"""

    qdb.sql_connection.perform_as_transaction(sql, [vals])

    return missing_files

def get_lat_longs():

    """Retrieve the latitude and longitude of all the public samples in the DB

    Returns

    -------

    list of [float, float]

        The latitude and longitude for each sample in the database

    """

    with qdb.sql_connection.TRN:

        # getting all the public studies

        studies = qdb.study.Study.get_by_status('public')

        results = []

        if studies:

            # we are going to create multiple union selects to retrieve the

            # latigute and longitude of all available studies. Note that

            # UNION in PostgreSQL automatically removes duplicates

            sql_query = """

                SELECT {0}, CAST(sample_values->>'latitude' AS FLOAT),

                       CAST(sample_values->>'longitude' AS FLOAT)

                FROM qiita.sample_{0}

                WHERE sample_values->>'latitude' != 'NaN' AND

                      sample_values->>'longitude' != 'NaN' AND

                      isnumeric(sample_values->>'latitude') AND

                      isnumeric(sample_values->>'longitude')"""

            sql = [sql_query.format(s.id) for s in studies]

            sql = ' UNION '.join(sql)

            qdb.sql_connection.TRN.add(sql)

            # note that we are returning set to remove duplicates

            results = qdb.sql_connection.TRN.execute_fetchindex()

        return results

def generate_biom_and_metadata_release(study_status='public'):

    """Generate a list of biom/meatadata filepaths and a tgz of those files

    Parameters

    ----------

    study_status : str, optional

        The study status to search for. Note that this should always be set

        to 'public' but having this exposed helps with testing. The other

        options are 'private' and 'sandbox'

    """

    studies = qdb.study.Study.get_by_status(study_status)

    qiita_config = ConfigurationManager()

    working_dir = qiita_config.working_dir

    portal = qiita_config.portal

    bdir = qdb.util.get_db_files_base_dir()

    time = datetime.now().strftime('%m-%d-%y %H:%M:%S')

    data = []

    for s in studies:

        # [0] latest is first, [1] only getting the filepath

        sample_fp = relpath(s.sample_template.get_filepaths()[0][1], bdir)

        for a in s.artifacts(artifact_type='BIOM'):

            if a.processing_parameters is None or a.visibility != study_status:

                continue

            merging_schemes, parent_softwares = a.merging_scheme

            software = a.processing_parameters.command.software

            software = '%s v%s' % (software.name, software.version)

            for x in a.filepaths:

                if x['fp_type'] != 'biom' or 'only-16s' in x['fp']:

                    continue

                fp = relpath(x['fp'], bdir)

                for pt in a.prep_templates:

                    categories = pt.categories

                    platform = ''

                    target_gene = ''

                    if 'platform' in categories:

                        platform = ', '.join(

                            set(pt.get_category('platform').values()))

                    if 'target_gene' in categories:

                        target_gene = ', '.join(

                            set(pt.get_category('target_gene').values()))

                    for _, prep_fp in pt.get_filepaths():

                        if 'qiime' not in prep_fp:

                            break

                    prep_fp = relpath(prep_fp, bdir)

                    # format: (biom_fp, sample_fp, prep_fp, qiita_artifact_id,

                    #          platform, target gene, merging schemes,

                    #          artifact software/version,

                    #          parent sofware/version)

                    data.append((fp, sample_fp, prep_fp, a.id, platform,

                                 target_gene, merging_schemes, software,

                                 parent_softwares))

    # writing text and tgz file

    ts = datetime.now().strftime('%m%d%y-%H%M%S')

    tgz_dir = join(working_dir, 'releases')

    create_nested_path(tgz_dir)

    tgz_name = join(tgz_dir, '%s-%s-building.tgz' % (portal, study_status))

    tgz_name_final = join(tgz_dir, '%s-%s.tgz' % (portal, study_status))

    txt_lines = [

        "biom fp\tsample fp\tprep fp\tqiita artifact id\tplatform\t"

        "target gene\tmerging scheme\tartifact software\tparent software"]

    with topen(tgz_name, "w|gz") as tgz:

        for biom_fp, sample_fp, prep_fp, aid, pform, tg, ms, asv, psv in data:

            txt_lines.append("%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (

                biom_fp, sample_fp, prep_fp, aid, pform, tg, ms, asv, psv))

            tgz.add(join(bdir, biom_fp), arcname=biom_fp, recursive=False)

            tgz.add(join(bdir, sample_fp), arcname=sample_fp, recursive=False)

            tgz.add(join(bdir, prep_fp), arcname=prep_fp, recursive=False)

        info = TarInfo(name='%s-%s-%s.txt' % (portal, study_status, ts))

        txt_hd = BytesIO()

        txt_hd.write(bytes('\n'.join(txt_lines), 'ascii'))

        txt_hd.seek(0)

        info.size = len(txt_hd.read())

        txt_hd.seek(0)

        tgz.addfile(tarinfo=info, fileobj=txt_hd)

    with open(tgz_name, "rb") as f:

        md5sum = md5()

        for c in iter(lambda: f.read(4096), b""):

            md5sum.update(c)

    move(tgz_name, tgz_name_final)

    vals = [

        ('filepath', tgz_name_final[len(working_dir):], r_client.set),

        ('md5sum', md5sum.hexdigest(), r_client.set),

        ('time', time, r_client.set)]

    for k, v, f in vals:

        redis_key = '%s:release:%s:%s' % (portal, study_status, k)

        # important to "flush" variables to avoid errors

        r_client.delete(redis_key)

        f(redis_key, v)

def generate_plugin_releases():

    """Generate releases for plugins

    """

    ARCHIVE = qdb.archive.Archive

    qiita_config = ConfigurationManager()

    working_dir = qiita_config.working_dir

    commands = [c for s in qdb.software.Software.iter(active=True)

                for c in s.commands if c.post_processing_cmd is not None]

    tnow = datetime.now()

    ts = tnow.strftime('%m%d%y-%H%M%S')

    tgz_dir = join(working_dir, 'releases', 'archive')

    create_nested_path(tgz_dir)

    tgz_dir_release = join(tgz_dir, ts)

    create_nested_path(tgz_dir_release)

    for cmd in commands:

        cmd_name = cmd.name

        mschemes = [v for _, v in ARCHIVE.merging_schemes().items()

                    if cmd_name in v]

        for ms in mschemes:

            ms_name = sub('[^0-9a-zA-Z]+', '', ms)

            ms_fp = join(tgz_dir_release, ms_name)

            create_nested_path(ms_fp)

            pfp = join(ms_fp, 'archive.json')

            archives = {k: loads(v)

                        for k, v in ARCHIVE.retrieve_feature_values(

                              archive_merging_scheme=ms).items()

                        if v != ''}

            with open(pfp, 'w') as f:

                dump(archives, f)

            # now let's run the post_processing_cmd

            ppc = cmd.post_processing_cmd

            # concatenate any other parameters into a string

            params = ' '.join(["%s=%s" % (k, v) for k, v in

                              ppc['script_params'].items()])

            # append archives file and output dir parameters

            params = ("%s --fp_archive=%s --output_dir=%s" % (

                params, pfp, ms_fp))

            ppc_cmd = "%s %s %s" % (

                ppc['script_env'], ppc['script_path'], params)

            p_out, p_err, rv = qdb.processing_job._system_call(ppc_cmd)

            p_out = p_out.rstrip()

            if rv != 0:

                raise ValueError('Error %d: %s' % (rv, p_out))

            p_out = loads(p_out)

    # tgz-ing all files

    tgz_name = join(tgz_dir, 'archive-%s-building.tgz' % ts)

    tgz_name_final = join(tgz_dir, 'archive.tgz')

    with topen(tgz_name, "w|gz") as tgz:

        tgz.add(tgz_dir_release, arcname=basename(tgz_dir_release))

    # getting the release md5

    with open(tgz_name, "rb") as f:

        md5sum = md5()

        for c in iter(lambda: f.read(4096), b""):

            md5sum.update(c)

    move(tgz_name, tgz_name_final)

    vals = [

        ('filepath', tgz_name_final[len(working_dir):], r_client.set),

        ('md5sum', md5sum.hexdigest(), r_client.set),

        ('time', tnow.strftime('%m-%d-%y %H:%M:%S'), r_client.set)]

    for k, v, f in vals:

        redis_key = 'release-archive:%s' % k

        # important to "flush" variables to avoid errors

        r_client.delete(redis_key)

        f(redis_key, v)

def get_software_commands(active):

    software_list = [s for s in qdb.software.Software.iter(active=active)]

    software_commands = defaultdict(lambda: defaultdict(list))

    for software in software_list:

        sname = software.name

        sversion = software.version

        commands = software.commands

        for command in commands:

            software_commands[sname][sversion].append(command.name)

        software_commands[sname] = dict(software_commands[sname])

    return dict(software_commands)

def update_resource_allocation_redis(active=True):

    """Updates redis with plots and information about current software.

    Parameters

    ----------

    active: boolean, optional

        Defaults to True. Should only be False when testing.

    """

    time = datetime.now().strftime('%m-%d-%y')

    scommands = get_software_commands(active)

    redis_key = 'resources:commands'

    r_client.set(redis_key, str(scommands))

    for sname, versions in scommands.items():

        for version, commands in versions.items():

            for cname in commands:

                col_name = "samples * columns"

                df = retrieve_resource_data(cname, sname, version, COLUMNS)

                if len(df) == 0:

                    continue

                fig, axs = resource_allocation_plot(df, col_name)

                titles = [0, 0]

                images = [0, 0]

                # Splitting 1 image plot into 2 separate for better layout.

                for i, ax in enumerate(axs):

                    titles[i] = ax.get_title()

                    ax.set_title("")

                    # new_fig, new_ax – copy with either only memory plot or

                    # only time

                    new_fig = plt.figure()

                    new_ax = new_fig.add_subplot(111)

                    line = ax.lines[0]

                    new_ax.plot(line.get_xdata(), line.get_ydata(),

                                linewidth=1, color='orange')

                    handles, labels = ax.get_legend_handles_labels()

                    for handle, label, scatter_data in zip(handles,

                                                           labels,

                                                           ax.collections):

                        color = handle.get_facecolor()

                        new_ax.scatter(scatter_data.get_offsets()[:, 0],

                                       scatter_data.get_offsets()[:, 1],

                                       s=scatter_data.get_sizes(), label=label,

                                       color=color)

                    new_ax.set_xscale('log')

                    new_ax.set_yscale('log')

                    new_ax.set_xlabel(ax.get_xlabel())

                    new_ax.set_ylabel(ax.get_ylabel())

                    new_ax.legend(loc='upper left')

                    new_fig.tight_layout()

                    plot = BytesIO()

                    new_fig.savefig(plot, format='png')

                    plot.seek(0)

                    img = 'data:image/png;base64,' + quote(

                          b64encode(plot.getvalue()).decode('ascii'))

                    images[i] = img

                    plt.close(new_fig)

                plt.close(fig)

                # SID, CID, col_name

                values = [

                    ("img_mem", images[0], r_client.set),

                    ("img_time", images[1], r_client.set),

                    ('time', time, r_client.set),

                    ("title_mem", titles[0], r_client.set),

                    ("title_time", titles[1], r_client.set)

                ]

                for k, v, f in values:

                    redis_key = 'resources$#%s$#%s$#%s$#%s:%s' % (

                                cname, sname, version, col_name, k)

                    r_client.delete(redis_key)

                    f(redis_key, v)