import torch
import os
import subprocess as sp
from mlutils.pt.training import TrainerBatch
from mlutils.callbacks import Callback
from os import path
import numpy as np
import json

class WeightedSum:
    def __init__(self, name, init, postprocessing=None):
        self.name = name
        self.v = init
        self.w = 0
        self.postprocessing = postprocessing

    def add(self, v, w):
        self.v = self.v + v * w
        self.w = self.w + w

    def get(self):
        # No accumulated value
        if self.w == 0:
            return 0
        v = self.v / self.w
        if self.postprocessing is not None:
            v = self.postprocessing(v)
        return v

    def __repr__(self):
        return '{} {:.5f}'.format(self.name, self.get())

class PerplexStatistics:

    def __init__(self):

        def _item(x):
            return x.item()

        def _exp_item(x):
            return torch.exp(x).item()

        self.stat = {
            'ppx': (WeightedSum('ppx', 0, _exp_item), '', ''),
            'ppx_doc': (WeightedSum('ppx_doc', 0, _exp_item), '', ''),
            'loss': (WeightedSum('loss', 0, _item), 'loss', 'doc_count'),
            'loss_rec': (WeightedSum('loss_rec', 0, _item), 'rec_loss', 'doc_count'),
            'kld': (WeightedSum('kld', 0, _item), 'kld', 'doc_count'),
            'penalty': (WeightedSum('penalty', 0, _item), 'penalty', 'doc_count'),
            'penalty_mean': (WeightedSum('penalty_mean', 0, _item), 'penalty_mean', 'doc_count'),
            'penalty_var': (WeightedSum('penalty_var', 0, _item), 'penalty_var', 'doc_count'),
        }

    def add(self, stat):
        """Accumulate statistics."""
        with torch.no_grad():
            data_batch = stat.pop('data')
            weight = {
                'word_count': data_batch.sum(),
                'doc_count': len(data_batch)
            }

            for s, k, w in self.stat.values():
                if s.name == 'ppx_doc':
                    s.add((stat['minus_elbo'] / data_batch.sum(dim=-1)).sum() / weight['doc_count'],
                          weight['doc_count'])
                elif s.name == 'ppx':
                    s.add(stat['minus_elbo'].sum() / weight['word_count'], weight['word_count'])
                else:
                    if k not in stat:  # skip for compatibility of multiple models.
                        continue
                    s.add(stat[k].mean(), weight[w])
        return self

    def description(self, prefix=''):
        return ' | '.join(['{} {:.5f}'.format(prefix + k, v)
                           for k, v in self.get_dict().items()])

    def get_value(self, k):
        """Get the accumulated value."""
        return self.stat[k][0].get()

    def get_dict(self):
        r = {}
        for k in self.stat.keys():
            t = self.stat[k][0].get()
            if t != 0:
                r[k] = t
        return r

class BatchOperation(TrainerBatch):

    def __init__(self, model, optimizer, loss, device, test_sample=1):
        self.model = model.cuda()
        self.optimizer = optimizer
        self.dst_device = device
        assert loss in ('mean', 'sum'), 'loss should be mean or sum of batch losses.'
        self.loss = loss
        self.test_sample = test_sample

    def train_batch(self, data, train=True):
        self.model.train(train)
        data = self.place_data_to_device(data)

        if train:
            # enter GSM class - forward function
            # out: 'stat' in GSM forward function; torch.Tensor
            out, topic = self.model(data)

            # self.optimizer.zero_grad()
            # if self.loss == 'mean':
            #     out['loss'].mean().backward()
            # else:
            #     out['loss'].sum().backward()
            # self.optimizer.step()
        else:
            with torch.no_grad():
                out, topic = self.model(data)

        out.update(data=data)
        return out, topic

def torch_detach(x):
    return x.detach().cpu().numpy()

def save_topics(save_path, vocab, topic_prob, topk=100, logger=None):
    """topic_prob: n_topic x vocab_size.
    Assumed that topic_prob[i] is probability distribution for all i.
    """
    if logger:
        logger.info('saving topics to {}'.format(save_path))

    values, indices = torch.topk(topic_prob, k=topk, dim=-1)
    indices = torch_detach(indices)
    values = torch_detach(values)

    topics = []
    for t in indices:
        topics.append(' '.join([vocab.itos[i] for i in t]))
    with open(save_path+'.topics', 'w') as f:
        f.write('\n'.join(topics))

    str_values = []
    for t in values:
        str_values.append(' '.join([str(v) for v in t]))
    with open(save_path+'.values', 'w') as f:
        f.write('\n'.join(str_values))

    torch.save(topic_prob, save_path + '.pt')

def evaluate_topic_coherence(topic_path, ref_corpus_dir, res_path, logger):
    """Evaluating topic coherence at topic_path whose lines are topics top 10 words.
    The evaluation uses the script at scripts/topics-20news.sh
    """

    if not os.path.exists(topic_path):
        logger.warning('topic file {} not exists'.format(topic_path))
        return -1

    v = -1
    try:
        p = sp.run(['bash', 'scripts/topic_coherence.sh', topic_path, ref_corpus_dir, res_path],
                   encoding='utf-8', timeout=20,
                   stdout=sp.PIPE, stderr=sp.DEVNULL)
        v = float(p.stdout.split('\n')[-3].split()[-1])
    except (ValueError, IndexError, TimeoutError):
        logger.warning('error when calculating topic coherence at {}'.format(topic_path))

    return v

def normalize(v2d, eps=1e-12):
    return v2d / (np.linalg.norm(v2d, axis=1, keepdims=True) + eps)

def wetc(e):
    """embedding matrix: N x D where N is the first N words in a topic, D is the embedding dimension."""
    e = normalize(e)
    t = normalize(e.mean(axis=0, keepdims=True))
    return float(e.dot(t.T).mean())

class EvaluationCallback(Callback):

    def __init__(self, base_dir, vocab, topk=10, corpus_dir="", embedding_path="", metric='npmi', every=10):
        """Evaluate topic coherence based on NPMI or WETC.
        For NPMI: args are, vocab, topk, corpus_dir
        For WETC: args are, embedding_path
        """
        super(EvaluationCallback, self).__init__()
        self.base_dir = base_dir
        self.vocab = vocab
        self.topk = topk
        self.corpus_dir = corpus_dir
        self.every = every
        self.cnt = 0
        self.max_tc = 0
        self.last_tc = 0
        if metric == 'wetc':
            assert os.path.exists(embedding_path), 'embedding file does not exists.'
            self.embedding = np.load(embedding_path)
            assert len(self.embedding) == len(vocab)
        metric = metric.lower()
        assert metric in ['npmi', 'wetc']
        self.metric = metric

    def wetc(self, topics):
        topics = torch_detach(topics)
        idx = np.argsort(topics, axis=1)
        tc = [wetc(self.embedding[idx[-self.topk:]]) for idx in idx]
        save_path = path.join(self.base_dir, 'wetc-topic-{}'.format(self.cnt))
        tc_mean = np.mean(tc)
        tc.append('mean {}'.format(tc_mean))
        with open(save_path, 'w') as f:
            json.dump(tc, f)
        return tc_mean

    def npmi(self, topics):
        save_path = path.join(self.base_dir, 'topic-{}'.format(self.cnt))
        save_topics(save_path, self.vocab, topics, self.topk,
                    self.trainer.logger)
        return evaluate_topic_coherence(save_path+'.topics', self.corpus_dir, save_path+'.res', self.trainer.logger)

    def evaluate_topic_coherence(self):
        topics = self.trainer.trainer_batch.model.get_topics()
        assert topics.size(1) == len(self.vocab), 'topics shape error, should be vocab size {}'.format(len(self.vocab))
        if self.metric == 'npmi':
            tc = self.npmi(topics)
        else:
            tc = self.wetc(topics)
        self.max_tc = max(self.max_tc, tc)
        self.last_tc = tc
        self.trainer.summary_writer.add_scalar('topic_coherence', tc, self.cnt)
        self.trainer.logger.info('topic coherence {}'.format(tc))

    def on_epoch_end(self, epoch, logs=None):
        self.cnt += 1
        if epoch % self.every == 0:
            self.evaluate_topic_coherence()

    def on_train_end(self, logs=None):
        self.evaluate_topic_coherence()

    def get_dict(self):
        return {
            'max_topic_coherence': self.max_tc,
            'last_topic_coherence': self.last_tc
        }

def recover_topic_embedding(topic_word_paths, embedding_path, dataset_dir):
    """Evaluate the WETC of topics generated by NPMI metric."""
    from data_utils import read_dataset
    assert isinstance(topic_word_paths, list), 'Multiple paths should be specified.'
    _, _, vocab = read_dataset(dataset_dir)
    embedding = np.load(embedding_path)
    scores = []
    for p in topic_word_paths:
        with open(p) as f:
            r = []
            for line in f:
                idx = [int(vocab.stoi[w]) for w in line.split()]
                r.append(wetc(embedding[idx]))
            scores.append(r)
    return np.array(scores)

def recover_model(exp_path):
    import json
    from mlutils.exp import yaml_load

    config_path = path.join(exp_path, 'trainer.json')
    if path.exists(config_path):
        print('find trainer.json')
        with open(config_path) as f:
            trainer = json.load(f)
            return BatchOperation.from_config(trainer['trainer_batch'])

    config_path = path.join(exp_path, 'trainer.yaml')
    if path.exists(config_path):
        print('find trainer.yaml')
        trainer = yaml_load(config_path)
        return BatchOperation.from_config(trainer['trainer_batch'])

    return None

def visualize_scale(count_path, scale_path, dataset_dir):
    from data_utils import read_dataset
    from mlutils.exp import yaml_load
    from matplotlib import pyplot as plt
    _, _, vocab = read_dataset(dataset_dir)
    scale = np.load(scale_path)
    if len(scale.shape) == 2:
        scale = scale[0]
    count = yaml_load(count_path)
    kv = sorted(count.items(), key=lambda x: -x[1])
    s = scale[[vocab.stoi[w[0]] for w in kv]]
    # scale = np.array([scale[vocab.stoi[w[0]]] for w in kv])
    plt.plot(s)


from astropy.io import fits
from astropy.table import Table
import argparse
import sys
import os
import numpy as np
from math import pi
import pandas
import copy

import arttools
from arttools._det_spatial import get_shadowed_pix_mask_for_urddata
from arttools.energy import get_events_energy
from arttools.caldb import get_shadowmask, get_energycal
from arttools.time import get_gti

from astropy.table import Table

import socket
import smtplib
from email.message import EmailMessage

parser = argparse.ArgumentParser(description="process L0 data to L1 format")
parser.add_argument("stem", help="part of the L0 files name, which are euqal to them")

ARTCALDBPATH = os.environ["ARTCALDB"]
indexfname = "artxc_index.fits"
#caldbindex = pandas.DataFrame(fits.getdata(indexfname, "CIF"))

URDTOTEL = {28: "T1",
            22: "T2",
            23: "T3",
            24: "T4",
            25: "T5",
            26: "T6",
            30: "T7"}

if __name__ == "__main__":
    try:
        sock = socket.socket()
        sock.connect(("10.5.2.24", 8081))
        sock.send(str.encode(sys.argv[1]))
    except ConnectionRefusedError as conerr:
        msg = EmailMessage()
        msg["to"] = "san@iki.rssi.ru"
        msg["from"] = "art@cosmos.ru"
        msg["subbj"] = "makeL1b fail to start daemon"
        msg.set_content(conerr)
        smtplib.SMTP("localhost").send_message(msg)
    finally:
        if len(sys.argv) != 4 or "-h" in sys.argv:
            print("description run like that 'python3 L1toL15.py stem outdir'"\
                    ", where stem is srg_20190727_214739_000")
            raise ValueError("wrong arguments")
        fname = sys.argv[1]
        stem = fname.rsplit(".")[0]
        outdir = sys.argv[2]
        attfname = sys.argv[3]
        if os.path.abspath(outdir) == os.path.abspath(os.path.dirname(stem)):
            raise ValueError("The L0 files will be overwriten")

        if not os.path.exists(outdir):
            os.mkdir(outdir)

        attdata = arttools.plot.get_attdata(attfname)
        urdfname = fname
        urdfile = fits.open(urdfname)
        urddata = Table(urdfile["EVENTS"].data).as_array()
        flag = np.ones(urddata.size, np.uint8)

        locgti = get_gti(urdfile, gtiextname="KVEA_GTI")
        locgti.merge_joint()

        caldbfile = get_energycal(urdfile)
        bkigti = arttools.time.make_bki_gti(urdfile)
        flag[(locgti & attdata.gti).mask_outofgti_times(urddata["TIME"])] = 0

        RA, DEC = np.empty(urddata.size, np.double), np.empty(urddata.size, np.double)
        attmask = attdata.gti.mask_outofgti_times(urddata["TIME"])
        RA[attmask], DEC[attmask] = arttools.orientation.get_photons_sky_coord(urddata[attmask], urdfile["EVENTS"].header["URDN"], attdata)

        ENERGY, xc, yc, grades = get_events_energy(urddata, urdfile["HK"].data, caldbfile)

        shadow = get_shadowmask(urdfile)
        maskshadow = get_shadowed_pix_mask_for_urddata(urddata, shadow)
        flag[np.logical_not(maskshadow)] = 2
        flag[np.any([urddata["RAW_X"] == 0, urddata["RAW_X"] == 47, urddata["RAW_Y"] == 0, urddata["RAW_Y"] == 47], axis=0)] = 3
        h = copy.copy(urdfile["EVENTS"].header)
        h.pop("NAXIS2")

        if bkigti.exposure > 0:
            bkimask = bkigti.mask_outofgti_times(urddata["TIME"])
            flag[bkimask] += 100

        cols = urdfile["EVENTS"].data.columns
        cols.add_col(fits.Column(name="ENERGY", array=ENERGY, format="1D", unit="keV"))
        cols.add_col(fits.Column(name="RA", array=np.copy(RA*180./pi), format="1D", unit="deg"))
        cols.add_col(fits.Column(name="DEC", array=np.copy(DEC*180./pi), format="1D", unit="deg"))
        cols.add_col(fits.Column(name="GRADE", array=grades, format="I"))
        cols.add_col(fits.Column(name="FLAG", array=flag, format="I"))
        cols["TIME_I"].array = urddata["TIME_I"]

        newurdtable = fits.BinTableHDU.from_columns(cols, header=h)

        newurdtable.name = "EVENTS"
        gtitable = fits.BinTableHDU(Table(locgti.arr, names=("START", "STOP")), header=urdfile["GTI"].header)
        newfile = fits.HDUList([urdfile[0], newurdtable, urdfile["HK"], gtitable])
        newfile.writeto(os.path.join(outdir, os.path.basename(urdfname)), overwrite=True)

        if not os.path.exists(os.path.join(outdir, os.path.basename(attfname))):
            attfile = fits.open(attfname)
            hdus = [attfile[0], ]
            for telescope in arttools.telescope.TELESCOPES:
                ra, dec, roll = arttools.orientation.quat_to_pol_and_roll(attdata(attdata.times)*arttools.caldb.ARTQUATS[telescope])
                hdus.append(fits.BinTableHDU(Table(np.array([attdata.times, ra*180/pi, dec*180/pi, roll*180/pi]).T, names=("TIME", "RA", "DEC", "ROLL")), header={"TELESCP":telescope}, name=telescope))
            fits.HDUList(hdus).writeto(os.path.join(outdir, os.path.basename(attfname)))


# -*- coding: utf-8 -*-
"""
Academy Color Encoding System - Input Transform
===============================================

Defines the *Academy Color Encoding System* (ACES) *Input Transform* utilities:

-   :func:`colour.sd_to_aces_relative_exposure_values`

See Also
--------
`RGB Colourspaces Jupyter Notebook
<http://nbviewer.jupyter.org/github/colour-science/colour-notebooks/\
blob/master/notebooks/models/rgb.ipynb>`_

References
----------
-   :cite:`Forsythe2018` : Forsythe, A. (2018). Private Discussion with
    Mansencal, T.
-   :cite:`TheAcademyofMotionPictureArtsandSciences2014q` : The Academy of
    Motion Picture Arts and Sciences, Science and Technology Council, & Academy
    Color Encoding System (ACES) Project Subcommittee. (2014). Technical
    Bulletin TB-2014-004 - Informative Notes on SMPTE ST 2065-1 - Academy Color
    Encoding Specification (ACES). Retrieved from
    https://github.com/ampas/aces-dev/tree/master/documents
-   :cite:`TheAcademyofMotionPictureArtsandSciences2014r` : The Academy of
    Motion Picture Arts and Sciences, Science and Technology Council, & Academy
    Color Encoding System (ACES) Project Subcommittee. (2014). Technical
    Bulletin TB-2014-012 - Academy Color Encoding System Version 1.0 Component
    Names. Retrieved from
    https://github.com/ampas/aces-dev/tree/master/documents
-   :cite:`TheAcademyofMotionPictureArtsandSciencese` : The Academy of Motion
    Picture Arts and Sciences, Science and Technology Council, & Academy Color
    Encoding System (ACES) Project Subcommittee. (n.d.). Academy Color Encoding
    System. Retrieved February 24, 2014, from
    http://www.oscars.org/science-technology/council/projects/aces.html
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.colorimetry import ILLUMINANTS_SDS, sd_to_XYZ
from colour.models import XYZ_to_xy
from colour.models.rgb import (ACES_2065_1_COLOURSPACE, ACES_RICD, RGB_to_XYZ,
                               XYZ_to_RGB, normalised_primary_matrix)
from colour.utilities import from_range_1, tsplit

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013-2019 - Colour Developers'
__license__ = 'New BSD License - https://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'FLARE_PERCENTAGE', 'S_FLARE_FACTOR', 'sd_to_aces_relative_exposure_values'
]

FLARE_PERCENTAGE = 0.00500
S_FLARE_FACTOR = 0.18000 / (0.18000 + FLARE_PERCENTAGE)

def sd_to_aces_relative_exposure_values(
        sd,
        illuminant=ILLUMINANTS_SDS['D65'],
        apply_chromatic_adaptation=False,
        chromatic_adaptation_transform='CAT02'):
    """
    Converts given spectral distribution to *ACES2065-1* colourspace relative
    exposure values.

    Parameters
    ----------
    sd : SpectralDistribution
        Spectral distribution.
    illuminant : SpectralDistribution, optional
        *Illuminant* spectral distribution.
    apply_chromatic_adaptation : bool, optional
        Whether to apply chromatic adaptation using given transform.
    chromatic_adaptation_transform : unicode, optional
        **{'CAT02', 'XYZ Scaling', 'Von Kries', 'Bradford', 'Sharp',
        'Fairchild', 'CMCCAT97', 'CMCCAT2000', 'CAT02_BRILL_CAT', 'Bianco',
        'Bianco PC'}**,
        *Chromatic adaptation* transform.

    Returns
    -------
    ndarray, (3,)
        *ACES2065-1* colourspace relative exposure values array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``XYZ``    | [0, 100]              | [0, 1]        |
    +------------+-----------------------+---------------+

    -   The chromatic adaptation method implemented here is a bit unusual
        as it involves building a new colourspace based on *ACES2065-1*
        colourspace primaries but using the whitepoint of the illuminant that
        the spectral distribution was measured under.

    References
    ----------
    :cite:`Forsythe2018`,
    :cite:`TheAcademyofMotionPictureArtsandSciences2014q`,
    :cite:`TheAcademyofMotionPictureArtsandSciences2014r`,
    :cite:`TheAcademyofMotionPictureArtsandSciencese`

    Examples
    --------
    >>> from colour import COLOURCHECKERS_SDS
    >>> sd = COLOURCHECKERS_SDS['ColorChecker N Ohta']['dark skin']
    >>> sd_to_aces_relative_exposure_values(sd)  # doctest: +ELLIPSIS
    array([ 0.1171785...,  0.0866347...,  0.0589707...])
    >>> sd_to_aces_relative_exposure_values(sd,
    ...     apply_chromatic_adaptation=True)  # doctest: +ELLIPSIS
    array([ 0.1180766...,  0.0869023...,  0.0589104...])
    """

    shape = ACES_RICD.shape
    if sd.shape != ACES_RICD.shape:
        sd = sd.copy().align(shape)

    if illuminant.shape != ACES_RICD.shape:
        illuminant = illuminant.copy().align(shape)

    s_v = sd.values
    i_v = illuminant.values

    r_bar, g_bar, b_bar = tsplit(ACES_RICD.values)

    def k(x, y):
        """
        Computes the :math:`K_r`, :math:`K_g` or :math:`K_b` scale factors.
        """

        return 1 / np.sum(x * y)

    k_r = k(i_v, r_bar)
    k_g = k(i_v, g_bar)
    k_b = k(i_v, b_bar)

    E_r = k_r * np.sum(i_v * s_v * r_bar)
    E_g = k_g * np.sum(i_v * s_v * g_bar)
    E_b = k_b * np.sum(i_v * s_v * b_bar)

    E_rgb = np.array([E_r, E_g, E_b])

    # Accounting for flare.
    E_rgb += FLARE_PERCENTAGE
    E_rgb *= S_FLARE_FACTOR

    if apply_chromatic_adaptation:
        xy = XYZ_to_xy(sd_to_XYZ(illuminant) / 100)
        NPM = normalised_primary_matrix(ACES_2065_1_COLOURSPACE.primaries, xy)
        XYZ = RGB_to_XYZ(E_rgb, xy, ACES_2065_1_COLOURSPACE.whitepoint, NPM,
                         chromatic_adaptation_transform)
        E_rgb = XYZ_to_RGB(XYZ, ACES_2065_1_COLOURSPACE.whitepoint,
                           ACES_2065_1_COLOURSPACE.whitepoint,
                           ACES_2065_1_COLOURSPACE.XYZ_to_RGB_matrix)

    return from_range_1(E_rgb)


import numpy as np
from math import ceil
from typing import List, Dict
from .base import BaseWriter

try:
    from osgeo import gdal
except:
    import gdal

class ClsWriter(BaseWriter):
    def __init__(self, config: Dict) -> None:
        super(ClsWriter, self).__init__(config, "tif")
        driver = gdal.GetDriverByName("GTiff")
        self.dst_ds = driver.Create(
            self.save_path, ceil(self.width / self.block_size[0]),
            ceil(self.height / self.block_size[1]), 1, gdal.GDT_UInt16)
        self.geotf = list(self.geotf)
        self.geotf[1] = self.geotf[1] * self.block_size[0]
        self.geotf[5] = self.geotf[5] * self.block_size[1]
        self.geotf = tuple(self.geotf)
        self.dst_ds.SetGeoTransform(self.geotf)
        self.dst_ds.SetProjection(self.proj)
        self.band = self.dst_ds.GetRasterBand(1)

    def write(self, block: np.ndarray, start: List[int]) -> None:
        xoff = start[0] / self.block_size[0]
        yoff = start[1] / self.block_size[1]
        self.band.WriteArray(block, int(xoff), int(yoff))
        self.dst_ds.FlushCache()

    def close(self) -> None:
        self.dst_ds = None


# Copyright (c) 2020-2021 The Center for Theoretical Biological Physics (CTBP) - Rice University and Northeastern University
# This file is from the OpenSMOG project, released under the MIT License.

R"""
The :class:`~.OpenSMOG` classes perform molecular dynamics using Structure-Based Models (SBM) for biomolecular simulations.
:class:`~.OpenSMOG` uses force fields generated by SMOG 2, and it allows the simulations of a wide variety of potential forms, including commonly employed C-alpha and all-atom variants.
Details about the default models in SMOG 2 can be found in the following resources:
    - **SMOG server**: https://smog-server.org/smog2/
    - **C-alpha**: Clementi, C., Nymeyer, H. and Onuchic, J.N., 2000. Topological and energetic factors: what determines the structural details of the transition state ensemble and “en-route” intermediates for protein folding? An investigation for small globular proteins. Journal of molecular biology, 298(5), pp.937-953.
    - **All-Atom**: Whitford, P.C., Noel, J.K., Gosavi, S., Schug, A., Sanbonmatsu, K.Y. and Onuchic, J.N., 2009. An all‐atom structure‐based potential for proteins: bridging minimal models with all‐atom empirical forcefields. Proteins: Structure, Function, and Bioinformatics, 75(2), pp.430-441.
"""

from simtk.openmm.app import *
from simtk.openmm import *
from simtk.unit import *
import os
import numpy as np
import xml.etree.ElementTree as ET
import warnings
from lxml import etree
import sys
from .OpenSMOG_Reporter import forcesReporter

class SBM:

    R"""
    The :class:`~.SBM` class performs Molecular dynamics simulations using structure-based (SMOG) models to investigate a broad range of biomolecular dynamics, including domain rearrangements in proteins, folding and ligand binding in RNA, and large-scale rearrangements in ribonucleoprotein assemblies. In its simplest form, a structure-based model defines a particular structure (usually obtained from X-ray, or NMR, methods) as the energetic global minimum. Find more information about SMOG models and OpenSMOG at http://smog-server.org

    The :class:`~.SBM` sets the environment to start the molecular dynamics simulations.

    Args:

        time_step (float, required):
            Simulation time step in units of :math:`\tau`.
        collision_rate (float, required):
            Friction/Damping constant in units of reciprocal time (:math:`1/\tau`).
        r_cutoff (float, required):
            Cutoff distance to consider non-bonded interactions in units of nanometers.
        pbc (boolean, optional):
            Turn PBC on/off. Default value: False
        temperature (float, required):
            Temperature in reduced units.
        name (str):
            Name used in the output files. (Default value: :code:`OpenSMOG`).
    """

    def __init__(self, time_step, collision_rate, r_cutoff, temperature, cmm = True, pbc = False, name = "OpenSMOG", warn = True):
        self.printHeader()
        print("\nFor more information, including descriptions of units and an example for how to launch a\nsimulation with OpenSMOG, run the help module. For example, if you named your simulation\nobject \"SMOGMODEL\", then issue the command:\n>SMOGMODEL.help()\n")
        self.name = name
        self.warn = warn
        self.dt = time_step * picoseconds

        self.gamma = collision_rate / picosecond
        self.rcutoff = r_cutoff * nanometers
        if self.warn:
            if not time_step in [0.0005, 0.002]:
                print('Note: The given time_step value is not the one usually employed in the SBM models. Make sure this value is correct. The suggested values are: time_step=0.0005 for C-alpha and time_step = 0.002 for All-Atoms.')
            if collision_rate != 1.0:
                print('Note: The given collision_rate value is not the one usually employed in the SBM models. Make sure this value is correct. The suggested value is: collision_rate=1.0.')
            if not r_cutoff in [1.1 ,0.65]:
                print('Note: The given r_cutoff value is not the one usually employed in the SBM models with OpenSMOG. Make sure this value is correct. The suggested values for r_cutoff are: 1.1 for the default C-alpha model and 0.65 for the all-atom model.')

        self.temperature = (temperature / 0.00831446261815) * kelvin
        self.forceApplied = False
        self.loaded = False
        self.folder = "."
        self.forceCount = 0
        self.pbc=pbc
        self.cmm=cmm
        self.nonbonded_present=False

    def help(self):
        R"""Prints information about using OpenSMOG.
        """
        print("""

#############################################################################

                       Information about using OpenSMOG

#############################################################################

                              Units in OpenSMOG

Time: When using OpenSMOG, it is important to consider the units. Since
OpenSMOG is built on OpenMM libraries, some unit labels can be misleading.
Specifically, SMOG models use reduced units. While it is possible to estimate
the effective timescales in these models, OpenMM reporters will label the
timescales as ns. In terms of accounting, 1 \"ns\" is equal to 1000 reduced
time units. For detailed discussions of units, one can consult any number of
recent papers by the SMOG team, such as Yang et al. JCP 2019, 151, 085102.

Temperature: When initializing your simulation with the SBM class, the
temperature is provided in reduced units, where a typical protein will fold
at around 1 reduced unit. However, similar to other MD packages, the
Boltzmann constant is built into the software.  Accordingly, the reported
temperature will be equal to (reduced temperature)*120.272. So, if you
initialize your system with a temperature of 0.5, the reported temperature
would read 0.5*120.272~60.136.

Energy: OpenSMOG will report all energies in reduced units. However, the
OpenMM libraries will label the energies as units of kJ/mole. This label
should be ignored, unless you calibrated your SMOG model to use this energy
scale.

                                Usage Example

While a broad range of simulation protocols may be used, here is a very basic
example for how to launch simulations using the OpenSMOG module with OpenMM.

>from OpenSMOG import SBM

Choose some basic runtime settings.  We will call our system 2ci2
>SMOGrun = SBM(name='2ci2', time_step=0.002, collision_rate=1.0, r_cutoff=1.2, temperature=0.5)

Select a platform and GPU IDs (if needed)
>SMOGrun.setup_openmm(platform='cuda',GPUindex='default')

Decide where to save your data (here, output_2ci2
>SMOGrun.saveFolder('output_2ci2')

You may optionally save some input file names to variables
>SMOG_grofile = '2ci2.gro'
>SMOG_topfile = '2ci2.top'
>SMOG_xmlfile = '2ci2.xml'

Load your force field data
>SMOGrun.loadSystem(Grofile=SMOG_grofile, Topfile=SMOG_topfile, Xmlfile=SMOG_xmlfile)

Create the simulation, and prepare to run
>SMOGrun.createSimulation()

Decide how frequently to save data
>SMOGrun.createReporters(trajectory=True, energies=True, energy_components=True, interval=10**3)

Launch the simulation
>SMOGrun.run(nsteps=10**6, report=True, interval=10**3)

For more information and help, see the OpenSMOG and SMOG 2 websites.
OpenSMOG: https://opensmog.readthedocs.io/en/latest/
SMOG 2: https://smog-server.org

If you have questions/suggestions, you can also email us at info@smog-server.org
""")

    def minimize(self,tolerance=1.0,maxIterations=0):
        R"""Wrapper for minimization simulation.

         Args:

            tolerance (float, required):
                Stopping criteria value between iteration. When the error between iteration is below this value, the minimization stops. (Default value: :code:`1.0`).
            maxIteration (int, required):
                Number of maximum steps to be performed in the minimization simulation. (Default value: :code:`1.0`).
        """
        self.simulation.minimizeEnergy(tolerance=tolerance,maxIterations=maxIterations)
        # it is very important that we reset the velocities. minimization warps the velocity values
        self.simulation.context.setVelocitiesToTemperature(self.temperature*kelvin)

    def setup_openmm(self, platform='opencl', precision='single', GPUindex='default', integrator="langevin"):

        R"""Sets up the parameters of the simulation OpenMM platform.

        Args:

            platform (str, optional):
                Platform to use in the simulations. Opitions are *CUDA*, *OpenCL*, *HIP*, *CPU*, *Reference*. (Default value: :code:`OpenCL`).
            precision (str, optional):
                Numerical precision type of the simulation. Options are *single*, *mixed*, *double*. (Default value: :code:`single`).  For details check the `OpenMM Documentation <http://docs.openmm.org/latest/developerguide/developer.html#numerical-precision>`__.
            GPUIndex (str, optional):
                Set of Platform device index IDs. Ex: 0,1,2 for the system to use the devices 0, 1 and 2. (Use only when GPU != default).
            integrator (str):
                Integrator to use in the simulations. Options are *langevin*,  *variableLangevin*, *verlet*, *variableVerlet* and, *brownian*. (Default value: :code:`langevin`).
        """

        precision = precision.lower()
        if precision not in ["mixed", "single", "double"]:
            raise ValueError("Precision must be mixed, single or double")

        properties = {}
        properties["Precision"] = precision
        if GPUindex.lower() != "default":
            properties["DeviceIndex"] = GPUindex

        self.properties = properties

        if platform.lower() == "opencl":
            platformObject = Platform.getPlatformByName('OpenCL')

        elif platform.lower() == "reference":
            platformObject = Platform.getPlatformByName('Reference')
            self.properties = {}

        elif platform.lower() == "cuda":
            platformObject = Platform.getPlatformByName('CUDA')

        elif platform.lower() == "cpu":
            platformObject = Platform.getPlatformByName('CPU')
            self.properties = {}

        elif platform.lower() == "hip":
            platformObject = Platform.getPlatformByName('HIP')

        else:
            self.exit("\n!!!! Unknown platform !!!!\n")

        self.platform = platformObject

        if integrator.lower() == "langevin":
            self.integrator = LangevinIntegrator(self.temperature,
                self.gamma, self.dt)
            self.integrator_type = integrator
        else:
            self.integrator = integrator
            self.integrator_type = "UserDefined"

        self.forceDict = {}
        self.forcesDict = {}

    def saveFolder(self, folder):

        R"""Sets the folder path to save data.

        Args:

            folder (str, optional):
                Folder path to save the simulation data. If the folder path does not exist, the function will create the directory.
        """

        if os.path.exists(folder) == False:
            os.mkdir(folder)
        self.folder = folder

    def loadSystem(self, Grofile, Topfile, Xmlfile):

        R"""Loads the input files in the OpenMM system platform. The input files are generated using SMOG2 software with the flag :code:`-OpenSMOG`. Details on how to create the files can be found in the `SMOG2 User Manual <https://smog-server.org/smog2/>`__.
        A tutorial on how to generate the inputs files for default all-atom and C-alpha models can be found `here <https://opensmog.readthedocs.io>`__.

        Args:

            Grofile (file, required):
                Initial structure for the MD simulations in *.gro* file format generated by SMOG2 software with the flag :code:`-OpenSMOG`.  (Default value: :code:`None`).
            Topfile (file, required):
                Topology *.top* file format generated by SMOG2 software with the flag :code:`-OpenSMOG`. The topology file lists the interactions between the system atoms except for the "Native Contacts" potential that is provided to OpenSMOG in a *.xml* file. (Default value: :code:`None`).
            Xmlfile (file, required):
                The *.xml* file that contains the all information that defines the "Contact" potential. The *.xml* file is generated by SMOG2 software with the flag :code:`-OpenSMOG`, which support custom potential functions. (Default value: :code:`None`).
        """
        def _checknames(f1,f2,f3):
            fn1 = os.path.basename(f1).rsplit('.', 1)[0]
            fn2 = os.path.basename(f2).rsplit('.', 1)[0]
            fn3 = os.path.basename(f3).rsplit('.', 1)[0]
            if fn1 == fn2 and fn1 ==fn3:
                return False
            else:
                return True

        if _checknames(Grofile, Topfile, Xmlfile) and  self.warn :
            warnings.warn('The Gro, Top and Xml files have different prefixes. Most people use the same name, so this may be a mistake.')

        self.inputNames = [Grofile, Topfile, Xmlfile]

        self._check_file(Grofile, '.gro')
        self.loadGro(Grofile)

        self._check_file(Topfile, '.top')
        self.loadTop(Topfile)

        self._check_file(Xmlfile, '.xml')
        self.loadXml(Xmlfile)

        print("Loaded force field and config files.")

    def _check_file(self, filename, ext):
        if not (filename.lower().endswith(ext)):
            raise ValueError('Wrong file extension: {} must to be {} extension'.format(filename, ext))

    def loadGro(self, Grofile):
        R"""Loads the  *.gro* file format in the OpenMM system platform. The input files are generated using SMOG2 software with the flag :code:`-OpenSMOG`. Details on how to create the files can be found in the `SMOG2 User Manual <https://smog-server.org/smog2/>`__.
        A tutorial on how to generate the inputs files for the default all-atom and C-alpha models can be found `here <https://opensmog.readthedocs.io>`__.

        Args:

            Grofile (file, required):
                Initial structure for the MD simulations in *.gro* file format generated by SMOG2 software with the flag :code:`-OpenSMOG`.  (Default value: :code:`None`).
        """
        self.Gro = GromacsGroFile(Grofile)

    def loadTop(self, Topfile):
        R"""Loads the  *.top* file format in the OpenMM system platform. The input files are generated using SMOG2 software with the flag :code:`-OpenSMOG`. Details on how to create the files can be found in the `SMOG2 User Manual <https://smog-server.org/smog2/>`__.
        A tutorial on how to generate the inputs files for the default all-atom and C-alpha models can be found `here <https://opensmog.readthedocs.io>`__.

        Args:

            Topfile (file, required):
                Topology *.top* file format generated by SMOG2 software with the flag :code:`-OpenSMOG`. The topology file defines the interactions between atoms, except for the "Native Contacts" potential that is provided to OpenSMOG in the form of a *.xml* file. (Default value: :code:`None`).
        """
        if self.pbc == True:
            print('This simulation will use periodic boundary conditions')
            self.Top = GromacsTopFile(Topfile,unitCellDimensions=self.Gro.getUnitCellDimensions())
            self.system = self.Top.createSystem(nonbondedMethod=CutoffPeriodic,nonbondedCutoff=self.rcutoff,removeCMMotion = self.cmm)
        else:
            print('This simulation will not use Periodic boundary conditions')
            self.Top = GromacsTopFile(Topfile)
            self.system = self.Top.createSystem(nonbondedMethod=CutoffNonPeriodic,nonbondedCutoff=self.rcutoff,removeCMMotion = self.cmm)
        nforces = len(self.system.getForces())
        for force_id, force in enumerate(self.system.getForces()):
            force.setForceGroup(force_id)
            self.forcesDict[force.__class__.__name__] = force
            self.forceCount +=1

    def _splitForces_contacts(self):
        #Contacts
        cont_data=self.data['contacts']
        n_forces =  len(cont_data[3])
        forces = {}
        for n in range(n_forces):
            forces[cont_data[3][n]] = [cont_data[0][n], cont_data[1][n], cont_data[2][n]]
        self.contacts = forces

    def _splitForces_nb(self):
        #Contacts
        nb_data=self.data['nonbond']
        n_forces =  len(nb_data[0])
        forces = {}
        for n in range(n_forces):
            forces[nb_data[0][n]] = [nb_data[1][n], nb_data[2][n], nb_data[3][n]]
        self.nonbond = forces

    def _customSmogForce(self, name, data):
        #first set the equation
        contacts_ff = CustomBondForce(data[0])

        #second set the number of variable
        for pars in data[1]:
            contacts_ff.addPerBondParameter(pars)

        #third, apply the bonds from each pair of atoms and the related variables.
        pars = [pars for pars in data[1]]

        for iteraction in data[2]:
            atom_index_i = int(iteraction['i'])-1
            atom_index_j = int(iteraction['j'])-1
            parameters = [float(iteraction[k]) for k in pars]

            contacts_ff.addBond(atom_index_i, atom_index_j, parameters)
        #forth, if the are global variables, add them to the force
        if self.constants_present==True:
            for const_key in self.data['constants']:
                contacts_ff.addGlobalParameter(const_key,self.data['constants'][const_key])
        self.forcesDict[name] =  contacts_ff
        contacts_ff.setForceGroup(self.forceCount)
        self.forceCount +=1

    def _customSmogForce_nb(self, name, data):
        #first set the equation
        nonbond_ff = CustomNonbondedForce(data[0])

        #Define per particle parameters
        nonbond_ff.addPerParticleParameter('q')
        nonbond_ff.addPerParticleParameter('type')

        #If the are global variables, add them to the force
        if self.constants_present==True:
            for const_key in self.data['constants']:
                nonbond_ff.addGlobalParameter(const_key,self.data['constants'][const_key])

        #Add cutoff as global parameter
        nonbond_ff.addGlobalParameter('r_c',self.rcutoff.value_in_unit(nanometer))

        #Load old nonbonded force for later use
        original_nonbonded=self.system.getForce(0)
        original_customnonbonded=self.system.getForce(1)

        #Get atoms types and number of types
        atom_types=[]
        for i in range(len(data[2])):
            atom_types.append(data[2][i]['type1'])
        self.atom_types=np.unique(np.array(atom_types))
        natom_types=len(self.atom_types)
        #Generate tables for each parameter
        tables={}
        for par in data[1]:
            tables[par]=np.ones([natom_types,natom_types])*np.nan

        #Fill tables with nonbond_param
        for nonbond_params in data[2]:
            #Get atoms id (name to number)
            type1=np.where(self.atom_types==nonbond_params['type1'])[0][0]
            type2=np.where(self.atom_types==nonbond_params['type2'])[0][0]
            for par in tables:
                tables[par][type1][type2]=nonbond_params[par]
                tables[par][type2][type1]=nonbond_params[par]

        #Generate Function from tables
        table_fun={}
        for par in data[1]:
            #Check none have nans
            if np.sum(tables[par]==np.nan)!=0:
                raise ValueError('Nonbonded force parameter: {:} is not defined for all atom interactions'.format(par))
            table_fun[par]=Discrete2DFunction(natom_types,natom_types,list(np.ravel(tables[par])))
            nonbond_ff.addTabulatedFunction(par,table_fun[par])
        #Get exceptions from topfile import
        for i in range(original_customnonbonded.getNumExclusions()):
            exclusion_id = original_customnonbonded.getExclusionParticles(i)
            nonbond_ff.addExclusion(exclusion_id[0],exclusion_id[1])
        ## ADD PARTICLES TO THE FORCE BASED ON SPECIFYING TYPE AND CHARGE
        ## CHARGES
        ## From nonbondedforce we get the charges for each atom
        atom_charges=[]
        ## Loop over every atom
        for i in range(self.system.getNumParticles()):
            atom_charge=original_nonbonded.getParticleParameters(i)[0]
            atom_charges.append(atom_charge.value_in_unit(constants.elementary_charge))
        ## ATOM TYPES
        ## From molecule information we get atom types
        atom_types=[]
        ## Get name and multiplicity of each molecule
        molecules_keys=[self.Top._molecules[i][0] for i in range(len(self.Top._molecules)) ]
        molecules_mul=[self.Top._molecules[i][1] for i in range(len(self.Top._molecules)) ]
        ## Loop over molecules
        for i in range(len(molecules_keys)):
            molecule=molecules_keys[i]
            mult=molecules_mul[i]
            # Loop over atoms in each molecule
            for atom in self.Top._moleculeTypes[molecule].atoms:
                # If atoms are repeated, then include the same entrie (CL,CL,CL,... or MG,MG,MG,...)
                for _ in range(mult):
                    atom_types.append(atom[1])
        for i in range(self.system.getNumParticles()):
            # GET ATOM TYPE
            at_type=np.where(atom_types[i]==self.atom_types)[0][0]
            # GET ATOM CHARGE
            charge=atom_charges[i]
            # ADD PARTICLE TO EACH FORCE WITH CORRESPONDING CHARGE AND TYPE
            nonbond_ff.addParticle([charge,at_type])
        #Set cutoff and nonbonded method
        if self.pbc == True:
            nonbond_ff.setNonbondedMethod(NonbondedForce.CutoffPeriodic)
        else:
            nonbond_ff.setNonbondedMethod(NonbondedForce.CutoffNonPeriodic)

        nonbond_ff.setCutoffDistance(self.rcutoff.value_in_unit(nanometer))

        self.forcesDict['Nonbonded'+str(name)] =  nonbond_ff
        nonbond_ff.setForceGroup(self.forceCount)
        self.forceCount +=1

    def loadXml(self, Xmlfile):
        R"""Loads the  *.xml* file format in the OpenMM system platform. The input files are generated using SMOG2 software with the flag :code:`-OpenSMOG`. Details on how to create the files can be found in the `SMOG2 User Manual <https://smog-server.org/smog2/>`__.
        A tutorial on how to generate the inputs files for default all-atom and C-alpha models can be found `here <https://opensmog.readthedocs.io>`__.

        Args:

            Xmlfile (file, required):
                The *.xml* file that contains all information that defines the "Contact" potentials. The *.xml* file is generated by SMOG2 software with the flag :code:`-OpenSMOG`, which support custom potential energy functions. (Default value: :code:`None`).
        """

        def validate(Xmlfile):
            path = "share/OpenSMOG.xsd"
            pt = os.path.dirname(os.path.realpath(__file__))
            filepath = os.path.join(pt,path)

            xmlschema_doc = etree.parse(filepath)
            xmlschema = etree.XMLSchema(xmlschema_doc)

            xml_doc = etree.parse(Xmlfile)

            result = xmlschema.validate(xml_doc)
            log = xmlschema.error_log
            return result,log

        def import_xml2OpenSMOG(file_xml):
            XML_potential = ET.parse(file_xml)
            root = XML_potential.getroot()
            xml_data={}

            ## Constants
            self.constants_present=False
            if root.find('constants') != None:
                self.constants_present=True
                constants={}
                constants_xml = root.find('constants')
                for const in constants_xml.iter('constant'):
                    constants[const.attrib['name']]=float(const.attrib['value'])
                xml_data['constants']=constants

            ## Contacts
            Force_Names=[]
            Expression=[]
            Parameters=[]
            Pairs=[]
            if root.find('contacts') == None:
                raise ValueError("No contacts were found in the XML file")
            contacts_xml=root.find('contacts')
            for i in range(len(contacts_xml)):
                for name in contacts_xml[i].iter('contacts_type'):
                    Force_Names.append(name.attrib['name'])

                for expr in contacts_xml[i].iter('expression'):
                    Expression.append(expr.attrib['expr'])

                internal_Param=[]
                for par in contacts_xml[i].iter('parameter'):
                    internal_Param.append(par.text)
                Parameters.append(internal_Param)

                internal_Pairs=[]
                for atompairs in contacts_xml[i].iter('interaction'):
                        internal_Pairs.append(atompairs.attrib)
                Pairs.append(internal_Pairs)

            xml_data['contacts']=[Expression,Parameters,Pairs,Force_Names]

            #Launch contact force function
            self.nonbond_present=False
            if root.find('nonbond') != None:
                self.nonbond_present=True
                nonbond_xml=root.find('nonbond')
                NonBond_Num=[]
                NBExpression=[]
                NBExpressionParameters=[]
                NBParameters=[]
                nonbond_xml=root.find('nonbond')
                for i in range(len(nonbond_xml)):
                    for types in nonbond_xml[i].iter('nonbond_bytype'):
                        NonBond_Num.append(i)
                    for expr in nonbond_xml[i].iter('expression'):
                        NBExpression.append(expr.attrib['expr'])
                    internal_Param=[]
                    for par in nonbond_xml[i].iter('parameter'):
                        internal_Param.append(par.text)
                    NBExpressionParameters.append(internal_Param)
                    internal_NBParam=[]
                    for nbpar in nonbond_xml[i].iter('nonbond_param'):
                        internal_NBParam.append(nbpar.attrib)
                    NBParameters.append(internal_NBParam)
                xml_data['nonbond']=[NonBond_Num,NBExpression,NBExpressionParameters,NBParameters]
            else:
                self.nonbond_present=False
            return xml_data

        if not (self.forceApplied):
            result,log=validate(Xmlfile)

            if not result:
                raise ValueError("The xml file is not adhere to the required schema (same as that used by SMOG 2). Error message:\n\n"+str(log)+"\n\n This typically means your xml file was corrupted, or you are using an incompatible version of SMOG 2")

            self.data = import_xml2OpenSMOG(Xmlfile)
            self._splitForces_contacts()
            for force in self.contacts:
                print("Creating Contacts force {:} from xml file".format(force))
                self._customSmogForce(force, self.contacts[force])
                self.system.addForce(self.forcesDict[force])

            if self.nonbond_present==True:
                self._splitForces_nb()
                for force in self.nonbond:
                    print("Creating Nonbonded force {:} from xml file".format(force))
                    self._customSmogForce_nb(force, self.nonbond[force])
                    self.system.addForce(self.forcesDict['Nonbonded'+str(force)])
                ## REMOVE OTHER NONBONDED FORCES
                self.system.removeForce(0)
                self.system.removeForce(0)

            self.forceApplied = True

        else:
            print('\n OpenSMOG forces already applied!!! \n')

    def createSimulation(self):

        R"""Creates the simulation context and loads into the OpenMM platform.
        """

        if not self.loaded:
            self.simulation = Simulation(self.Top.topology, self.system, self.integrator, self.platform, self.properties)
            self.simulation.context.setPositions(self.Gro.positions)
            self.simulation.context.setVelocitiesToTemperature(self.temperature)
            self.loaded = True
        else:
            print('\n Simulations context already created! \n')

    def _checkFile(self,filename):
        if os.path.isfile(filename):
            i = 1
            ck = True
            while i <= 10 and ck:
                newname = filename + "_" + str(i)
                if not os.path.isfile(newname):
                    print("{:} already exists.  Backing up to {:}".format(filename,newname))
                    os.rename(filename, newname)
                    ck = False
                else:
                    i += 1

    def createReporters(self, trajectory=True, trajectoryName=None, trajectoryFormat='dcd', energies=True, energiesName=None, energy_components=False, energy_componentsName=None, interval=1000):
        R"""Creates the reporters to provide the output data.

        Args:

            trajectory (bool, optional):
                 Whether to save the trajectory *.dcd* file containing the position of the atoms as a function of time. (Default value: :code:`True`).
            trajectoryName (str, optional):
                Name of the trajectory file.
            trajectoryFormat (str, optional):
                File format of the trajectory file. Options are dcd, pdb, pdbx, hdf5, netcdf, and xtc. Saving the trajectory in the file formats hdf5, netcdf, and xtc require the MDTraj package. (Default value: :code:`dcd`).
            energies (bool, optional):
                 Whether to save the energies in a *.txt* file containing five columns, comma-delimited. The header of the files shows the information of each collum: #"Step","Potential Energy (kJ/mole)","Kinetic Energy (kJ/mole)","Total Energy (kJ/mole)","Temperature (K)". (Default value: :code:`True`).
            energiesName (str, optional):
                Name of the energy file.
            energy_components (bool, optional):
                 Whether to save the potential energy for each applied force in a *.txt* file containing several columns, comma-delimited. The header of the files shows the information of each column. An example of the header is: #"Step","electrostatic","Non-Contacts","Bonds","Angles","Dihedrals","contact_1-10-12". (Default value: :code:`False`).
            energy_componentsName (str, optional):
                Name of the energy_components file.
            interval (int, required):
                 Frequency to write the data to the output files. (Default value: :code:`10**3`)
        """

        self.outputNames = []
        if trajectory:
            trajectoryFormat=trajectoryFormat.lower()
            if trajectoryName is None:
                trajfile = os.path.join(self.folder, self.name + '_trajectory.'+trajectoryFormat)
            else:
                trajfile = os.path.join(self.folder, trajectoryName + "."+trajectoryFormat)
            self._checkFile(trajfile)
            self.outputNames.append(trajfile)

            for i in ['hdf5', 'xtc', 'netcdf']:
                if i == trajectoryFormat:
                    print("""
The """+str(trajectoryFormat)+""" trajectory format requires mdtraj to be loaded.
Will try to import mdtraj...""")
                    import mdtraj as md

            if trajectoryFormat == 'dcd':
                self.simulation.reporters.append(DCDReporter(trajfile, interval))
            elif trajectoryFormat == 'pdb':
                self.simulation.reporters.append(PDBReporter(trajfile, interval))
            elif trajectoryFormat == 'pdbx':
                self.simulation.reporters.append(PDBxReporter(trajfile, interval))
            elif trajectoryFormat == 'hdf5':
                self.simulation.reporters.append(md.reporters.HDF5Reporter(trajfile, interval))
            elif trajectoryFormat == 'netcdf':
                self.simulation.reporters.append(md.reporters.NetCDFReporter(trajfile, interval))
            elif trajectoryFormat == 'xtc':
                self.simulation.reporters.append(md.reporters.XTCReporter(trajfile, interval))
            else:
                raise ValueError("Trajectory format "+str(trajectoryFormat)+" not recognized")

        if energies:
            if energiesName is None:
                energyfile = os.path.join(self.folder, self.name+ '_energies.txt')
            else:
                 energyfile = os.path.join(self.folder, energiesName + ".txt")
            self._checkFile(energyfile)
            self.outputNames.append(energyfile)
            self.simulation.reporters.append(StateDataReporter(energyfile, interval, step=True,
                                                          potentialEnergy=True, kineticEnergy=True,
                                                            totalEnergy=True,temperature=True, separator=","))

        if energy_components:
            if energy_componentsName is None:
                forcefile = os.path.join(self.folder, self.name + '_forces.txt')
            else:
                forcefile = os.path.join(self.folder, energy_componentsName + '.txt')
            self._checkFile(forcefile)
            self.outputNames.append(forcefile)
            self.simulation.reporters.append(forcesReporter(forcefile, interval, self.forcesDict, step=True))

    def run(self, nsteps, report=True, interval=10**4):
        R"""Run the molecular dynamics simulation.

        Args:

            nsteps (int, required):
                 Number of steps to be performed in the simulation. (Default value: :code:`10**7`)
            report (bool, optional):
                Whether to print the simulation progress. (Default value: :code:`True`).
            interval (int, required):
                 Frequency to print the simulation progress. (Default value: :code:`10**4`)
        """

        if report:
            self.simulation.reporters.append(StateDataReporter(sys.stdout, interval, step=True, remainingTime=True,
                                                  progress=True, speed=True, totalSteps=nsteps, separator="\t"))
        self._createLogfile()
        self.simulation.step(nsteps)

    def _createLogfile(self):
        import platform
        import datetime
        logFilename = os.path.join(self.folder, 'OpenSMOG.log')
        self._checkFile(logFilename)
        self.outputNames.append(logFilename)
        with open(logFilename, 'w') as f:
            ori = sys.stdout
            sys.stdout = f
            self.printHeader()
            sys.stdout = ori

        #system_information
            f.write('\nSystem Information:\n')
            f.write('-------------------\n')

            f.write('Date and time: {:}\n'.format(datetime.datetime.now()))
            f.write('Machine information: {:} : {:}, {:} : {:}\n'.format("System", platform.uname().system, "Version", platform.uname().version))
            f.write('Platform: {:}\n'.format(self.platform.getName()))
            if (self.platform.getName() in ["CUDA", "OpenCL", "HIP"]):
                f.write('Precision: {:}\n'.format(self.properties['Precision']))
            f.write('Integrator: {:}\n'.format(self.integrator_type))
            f.write('Savefolder: {:}\n'.format(self.folder))

            f.write('\nSimulation Information:\n')
            f.write('-----------------------\n')

            f.write('Name: {:}\n'.format(self.name))
            f.write('Time step: {:}\n'.format(self.dt/picoseconds))
            f.write('Collision Rate: {:}\n'.format(self.gamma*picosecond))
            f.write('r_Cutoff: {:}\n'.format(self.rcutoff/nanometers))
            f.write('Temperature: {:}\n'.format(self.temperature * 0.008314/kelvin))

            f.write('\nInput Files:\n')
            f.write('------------\n')
            f.write('GroFile: {:}\n'.format(self.inputNames[0]))
            f.write('TopFile: {:}\n'.format(self.inputNames[1]))
            f.write('XmlFile: {:}\n'.format(self.inputNames[2]))

            f.write('\nOutput Files:\n')
            f.write('-------------\n')
            for n in self.outputNames:
                f.write(n+"\n")

            sys.stdout = ori

    def printHeader(self):
        print('{:^96s}'.format("****************************************************************************************"))
        print('{:^96s}'.format("**** *** *** *** *** *** *** *** OpenSMOG-1.1.0 *** *** *** *** *** *** *** ****"))
        print('')
        print('{:^96s}'.format("The OpenSMOG classes perform molecular dynamics simulations using"))
        print('{:^96s}'.format("Structure-Based Models (SBM) for biomolecular systems,"))
        print('{:^96s}'.format("and it allows for the simulation of a wide variety of potential forms."))
        print('{:^96s}'.format("OpenSMOG uses force field files generated by SMOG 2."))
        print('{:^96s}'.format("OpenSMOG documentation is available at https://opensmog.readthedocs.io"))
        print('')
        print('{:^96s}'.format("OpenSMOG is described in: Oliveira and Contessoto et al,"))
        print('{:^96s}'.format("SMOG 2 and OpenSMOG: Extending the limits of structure-based models."))
        print('{:^96s}'.format("Protein Science, DOI:10.1002/pro.4209, 2021."))
        print('')
        print('{:^96s}'.format("This package is the product of contributions from a number of people, including:"))
        print('{:^96s}'.format("Jeffrey Noel, Mariana Levi, Antonio Oliveira, Vinícius Contessoto,"))
        print('{:^96s}'.format("Esteban Dodero-Rojas, Mohit Raghunathan, Joyce Yang, Prasad Bandarkar,"))
        print('{:^96s}'.format("Udayan Mohanty, Ailun Wang, Heiko Lammert, Ryan Hayes,"))
        print('{:^96s}'.format("Jose Onuchic & Paul Whitford"))
        print('')
        print('{:^96s}'.format("Copyright (c) 2021, The SMOG development team at"))
        print('{:^96s}'.format("Rice University and Northeastern University"))
        print('{:^96s}'.format("****************************************************************************************"))
        sys.stdout.flush()


import numpy as np
from astropy.wcs import WCS
from astropy.table import Table
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.io import fits
import sys

import urllib
from .mast import *
from .utils import *

__all__ = ['Source', 'multi_sectors']

def multi_sectors(sectors, tic=None, gaia=None, coords=None):
    """Obtain a list of Source objects for a single target, for each of multiple sectors for which the target was observed.

    Parameters
    ----------
    tic : int, optional
        The TIC ID of the source.
    gaia : int, optional
        The Gaia DR2 source_id.
    coords : tuple, optional
        The (RA, Dec) coords of the object in degrees.
    sectors : list or str
        The list of sectors for which data should be returned, or `all` to return all sectors
        for which there are data.
    """
    objs = []
    if sectors == 'all':
        sectors = list(np.arange(1,14,1, dtype=int))
    if type(sectors) == list:
        for s in sectors:
            star = Source(tic=tic, gaia=gaia, coords=coords, sector=int(s))
            if star.sector is not None:
                objs.append(star)
        return objs
    else:
        raise TypeError("Sectors needs to be either 'all' or a type(list) to work.")

def load_postcard_guide():
    """Load and return the postcard coordinates guide."""
    guide_link = urllib.request.urlopen('http://astro.uchicago.edu/~bmontet/TESS_postcards/postcard.guide')
    guide = guide_link.read().decode('utf-8')
    guide = Table.read(guide, format='ascii.basic') # guide to postcard locations
    return guide

class Source(object):
    """A single source observed by TESS.

    Parameters
    ----------
    tic : int, optional
        The TIC ID of the source.
    gaia : int, optional
        The Gaia DR2 source_id.
    coords : tuple or astropy.coordinates.SkyCoord, optional
        The (RA, Dec) coords of the object in degrees or an astropy SkyCoord object.
    fn : str, optional
        Filename of a TPF corresponding to the desired source.
    sector : int or str
        The sector for which data should be returned, or `recent` to
        obtain data for the most recent sector which contains this target.

    Attributes
    ----------
    tess_mag : float
        The TESS magnitude from the TIC.
    sector : int
        Sector in which source was observed by TESS.
    camera : int
        TESS camera on which source falls.
    chip : int
        TESS chip on which source falls.
    position_on_chip : (int, int)
        Predicted (x,y) coords of object center on chip.
    postcard : str
        Name of the best postcard (postcard where source should be located closest to the center).
    position_on_postcard : (int, int)
        Predicted (x, y) coords of object center on the above-named postcard.
        Does NOT take into account additional pointing corrections; these will be applied when making the TPF.
    all_postcards : list of strs
        Names of all postcards where the source appears.
    """
    def __init__(self, tic=None, gaia=None, coords=None, fn=None, sector=None):
        self.tic     = tic
        self.gaia    = gaia
        self.coords  = coords
        self.fn      = fn
        self.premade = False
        self.usr_sec = sector

        if self.fn is not None:
            try:
                hdu = fits.open(self.fn)
            except:
                assert False, "{0} is not a valid filename.".format(fn)
            hdr = hdu[0].header
            self.tic      = hdr['TIC_ID']
            self.tess_mag = hdr['TMAG']
            self.gaia     = hdr['GAIA_ID']
            self.coords   = (hdr['CEN_RA'], hdr['CEN_DEC'])
            self.premade  = True
            self.sector   = hdr['SECTOR']
            self.camera   = hdr['CAMERA']
            self.chip     = hdr['CHIP']
            self.position_on_chip = (hdr['CHIPPOS1'], hdr['CHIPPOS2'])
            self.position_on_postcard = (hdr['POSTPOS1'], hdr['POSTPOS2'])

        else:
            if self.coords is not None:
                if type(self.coords) is SkyCoord:
                    self.coords = (self.coords.ra.degree, self.coords.dec.degree)
                elif (len(self.coords) == 2) & all(isinstance(c, float) for c in self.coords):
                    self.coords  = coords
                else:
                    assert False, ("Source: invalid coords. Valid input types are: "
                                   "(RA [deg], Dec [deg]) tuple or astropy.coordinates.SkyCoord object.")

                self.tic, self.tess_mag, sep = tic_from_coords(self.coords)
                self.gaia = gaia_from_coords(self.coords)

            elif self.gaia is not None:
                self.coords = coords_from_gaia(self.gaia)
                self.tic, self.tess_mag = tic_from_coords(self.coords)

            elif self.tic is not None:
                self.coords, self.tess_mag = coords_from_tic(self.tic)
                self.gaia = gaia_from_coords(self.coords)

            else:
                assert False, ("Source: one of the following keywords must be given: "
                               "tic, gaia, coords, fn.")

            self.tess_mag = self.tess_mag[0]
            self.locate_on_tess() # sets sector, camera, chip, postcard,
                                  # position_on_chip, position_on_postcard

    def locate_on_chip(self, guide):
        """Finds the TESS sector, camera, chip, and position on chip for the source.
        Sets attributes sector, camera, chip, position_on_chip.
        """
        def cam_chip_loop(sec):
            for cam in np.unique(guide['CAMERA']):
                for chip in np.unique(guide['CCD']):
                    mask = ((guide['SECTOR'] == sec) & (guide['CAMERA'] == cam)
                            & (guide['CCD'] == chip))
                    if np.sum(mask) == 0:
                        continue # this camera-chip combo is not in the postcard database yet
                    else:
                        random_postcard = guide[mask][0]
                    d = {}
                    for j,col in enumerate(random_postcard.colnames):
                        d[col] = random_postcard[j]
                    hdr = fits.Header(cards=d) # make WCS info from one postcard header
                    xy = WCS(hdr).all_world2pix(self.coords[0], self.coords[1], 1, quiet=True) # position in pixels in FFI dims
                    x_zero, y_zero = hdr['POSTPIX1'], hdr['POSTPIX2']
                    xy = np.array([xy[0]+x_zero, xy[1]+y_zero])
                    if (0 <= xy[0] < 2048) & (44 <= xy[1] < 2092):
                        self.sector = sec
                        self.camera = cam
                        self.chip = chip
                        self.position_on_chip = np.ravel(xy)

        guide = load_postcard_guide()
        self.sector=None

        if self.usr_sec is None:
            for sec in np.unique(guide['SECTOR']):
                cam_chip_loop(sec)

            if self.sector is None:
                raise SearchError("TESS has not (yet) observed your target.")

        elif self.usr_sec is not None:
            if type(self.usr_sec) == int:
                cam_chip_loop(self.usr_sec)

            # Searches for the most recent sector the object was observed in
            elif self.usr_sec.lower() == 'recent':
                for s in np.arange(15,0,-1):
                    cam_chip_loop(s)
                    if self.sector is not None:
                        break

        return

    def locate_on_tess(self):
        """Finds the best TESS postcard(s) and the position of the source on postcard.
        Sets attributes postcard, position_on_postcard, all_postcards.
        """
        guide = load_postcard_guide()
        self.locate_on_chip(guide)

        if self.sector is None:
            return

        # Searches through postcards for the given sector, camera, chip
        in_file, dists=[], []

        # Searches through rows of the table
        for i in range(len(guide)):
            postcard_inds = np.arange(len(guide))[(guide['SECTOR'] == self.sector) & (guide['CAMERA'] == self.camera)
                                                  & (guide['CCD'] == self.chip)]
        xy = self.position_on_chip

        # Finds postcards containing the source
        for i in postcard_inds: # loop rows

            x_cen, y_cen= guide['CEN_X'][i], guide['CEN_Y'][i]
            l, w = guide['POST_H'][i]/2., guide['POST_W'][i]/2.

            # Checks to see if xy coordinates of source falls within postcard
            if (xy[0] >= x_cen-l) & (xy[0] <= x_cen+l) & (xy[1] >= y_cen-w) & (xy[1] <= y_cen+w):
                in_file.append(i)
                dists.append(np.min([xy[0]-(x_cen-l), (x_cen+l)-xy[0], xy[1]-(y_cen-w), (y_cen+w)-xy[1]]))

        # If more than one postcard is found for a single source, choose the postcard where the
        # source is closer to the center
        if in_file == []:
            raise SearchError("Sorry! We don't have a postcard for you at the moment.")
        elif len(in_file) > 1:
            best_ind = np.argmax(dists)
        else:
            best_ind = 0

        self.all_postcards = guide['POSTNAME'][in_file]
        self.postcard = guide['POSTNAME'][in_file[best_ind]]

        self.sector = guide['SECTOR'][in_file[best_ind]] # WILL BREAK FOR MULTI-SECTOR TARGETS
        self.camera = guide['CAMERA'][in_file[best_ind]]
        self.chip = guide['CCD'][in_file[best_ind]]

        i = in_file[best_ind]
        postcard_pos_on_ffi = (guide['CEN_X'][i] - guide['POST_H'][i]/2.,
                                guide['CEN_Y'][i] - guide['POST_W'][i]/2.)
        self.position_on_postcard = xy - postcard_pos_on_ffi # as accurate as FFI WCS

        i = in_file[best_ind]
        postcard_pos_on_ffi = (guide['CEN_X'][i] - guide['POST_H'][i]/2.,
                              guide['CEN_Y'][i] - guide['POST_W'][i]/2.) # (x,y) on FFI where postcard (x,y) = (0,0)
        self.position_on_postcard = xy - postcard_pos_on_ffi # as accurate as FFI WCS


'''
Clase que maneja la solicitud y emisión de recomendaciones
'''
from entities.recommendation import Recommendation
from entities.profile import Profile
from entities.automobile import Automobile
from clusteringModel.kmodesManager import KmodesManager
from recommenderCore.contentBased import ContentBased
from entities.opinionSheet import OpinionSheet
import numpy as np

class RecommendationManager:

    @staticmethod
    def getRecommendation(form,idReq,MyConnection):
        RecommendationManager.setResults(form, MyConnection,idReq)#1. Almaceno formulrio en base de datos.
        array=RecommendationManager.getNumpyForm(form)# Traslado formulario a un numerico numerico
        cluster=KmodesManager.getCluster(array,MyConnection)#2. Clasifico formulario para obtener perfil.
        idModel=MyConnection.getLastModel()[0]#obtengo el ultimo modelo definido (la posicion 0 indica el id)
        profile=MyConnection.getPerfil(cluster,idModel)# obtengo los datos de puntuaci'on de perfil
        idRecom=MyConnection.addRecom(idReq,profile[0]) #3. creo una recomendacion en base de datos y asigno el perfil(posicion 0 es el id del perfil)
        if(idRecom):#4.genero recomendacion %% requiere generate OVERVIEW
            print('::OK::Formulario recibido: ',array[0])

            '''
            autos1=ContentBased.getBestRatedAutos(False,cluster,idModel,MyConnection,55)# En caso de tener restricciones de autos, agregarlo como primer parametro, en caso contraario se introduce False
            print('Autos después del filtrado de perfil: ',autos1)
            autos1=ContentBased.getSimilarAutos(MyConnection,array[0],7,autos1) ## En caso de tener una resticccioon sobre ciertos automoviles, agregarlo como tercer parametro
            print('filtro basado en contenido ok: ',autos1)
            '''
            autos1=ContentBased.getSimilarAutos(MyConnection,array[0],25) ## En caso de tener una resticccioon sobre ciertos automoviles, agregarlo como segundo parametro
            #print('Autos después del filtrado de caracteristicas: ',autos1)
            autos1=ContentBased.getBestRatedAutos(autos1,cluster,idModel,MyConnection,7)# En caso de tener restricciones de autos, agregarlo como primer parametro, en caso contraario se introduce False
            print('Puntuaciones de autos: ',autos1)
            #'''

            autos1=ContentBased.getRestrictedAutos(MyConnection,autos1,Nresults=5,MaxMarca=3,Maxmodel=1)# Restricciones de marca y modelo por recomendación
            print('Resultados del usuario : ',autos1)
            j=1 # Se guardan los autos resultado de recomendación en base de datos
            for auto in autos1:
                if not MyConnection.addResultRecom(idRecom,j,auto+1):
                    print('error')
                    return False
                j+=1
            print('::OK:: Datos almacenados en BD')
            data_Autos=MyConnection.getAutosByIdRecom(idRecom[0])# 7. recupero autos para crear el objeto respuesta
            if(data_Autos):
                profileResponse=Profile(profile[0],profile[1],profile[2])
                arrAutosResponse=[]
                for data_Auto in data_Autos:
                    auxAuto=Automobile(data_Auto[1],data_Auto[2],data_Auto[3],data_Auto[4],data_Auto[5],)
                    arrAutosResponse.append(auxAuto)
                recomResponse=Recommendation(idRecom[0],arrAutosResponse,profileResponse)
                print(recomResponse.get_recommendation())
                return recomResponse.get_recommendation()
            else:
                return False
        else:
            print('Error al crear nueva recomendacion')
            return False

    # ----------------------- METODOS AUXILIARES ------------------
    @staticmethod
    def setResults(questions,MyConnection,idReq):
        for x in questions['questions']:
            idResp=x['answer']
            idQues=x['id']
            MyConnection.addResult(idReq,idQues,idResp)
        print('formulario guardado en DB')
        return True

    @staticmethod
    def getNumpyForm(questions):
        formList=[]
        for x in questions['questions']:
            formList.append(x['answer'])
        return np.asarray(formList).reshape(1, -1)#forma admitida


__author__ = "liuwei"

"""
A char word model
"""
import numpy as np
import torch
import torch.nn as nn

from model.cw_ner.modules.gaz_embed import Gaz_Embed
from model.cw_ner.modules.gaz_bilstm import Gaz_BiLSTM
from model.cw_ner.model.crf import CRF
from model.cw_ner.functions.utils import random_embedding
from model.cw_ner.functions.gaz_opt import get_batch_gaz
from model.cw_ner.functions.utils import reverse_padded_sequence

class CW_NER(torch.nn.Module):
    def __init__(self, data, type=1):
        print("Build char-word based NER Task...")
        super(CW_NER, self).__init__()

        self.gpu = data.HP_gpu
        self.label_size = data.label_alphabet_size
        self.type = type
        self.gaz_embed = Gaz_Embed(data, type)

        self.char_embedding = nn.Embedding(data.char_alphabet.size(), data.char_emb_dim)

        self.lstm = Gaz_BiLSTM(data, data.char_emb_dim + data.gaz_emb_dim, data.HP_hidden_dim)

        self.crf = CRF(data.label_alphabet_size, self.gpu)

        self.hidden2tag = nn.Linear(data.HP_hidden_dim * 2, data.label_alphabet_size + 2)

        if data.pretrain_char_embedding is not None:
            self.char_embedding.weight.data.copy_(torch.from_numpy(data.pretrain_char_embedding))
        else:
            self.char_embedding.weight.data.copy_(
                random_embedding(data.char_alphabet_size, data.char_emb_dim)
            )

        if self.gpu:
            self.char_embedding = self.char_embedding.cuda()
            self.hidden2tag = self.hidden2tag.cuda()

    def neg_log_likelihood_loss(self, gaz_list, reverse_gaz_list, char_inputs, char_seq_lengths, batch_label, mask):
        """
        get the neg_log_likelihood_loss
        Args:
            gaz_list: the batch data's gaz, for every chinese char
            reverse_gaz_list: the reverse list
            char_inputs: word input ids, [batch_size, seq_len]
            char_seq_lengths: [batch_size]
            batch_label: [batch_size, seq_len]
            mask: [batch_size, seq_len]
        """
        batch_size = char_inputs.size(0)
        seq_len = char_inputs.size(1)
        lengths = list(map(int, char_seq_lengths))

        # print('one ', reverse_gaz_list[0][:10])

        # get batch gaz ids
        batch_gaz_ids, batch_gaz_length, batch_gaz_mask = get_batch_gaz(reverse_gaz_list, batch_size, seq_len, self.gpu)

        # print('two ', batch_gaz_ids[0][:10])

        reverse_batch_gaz_ids, reverse_batch_gaz_length, reverse_batch_gaz_mask = get_batch_gaz(gaz_list, batch_size, seq_len, self.gpu)
        reverse_batch_gaz_ids = reverse_padded_sequence(reverse_batch_gaz_ids, lengths)
        reverse_batch_gaz_length = reverse_padded_sequence(reverse_batch_gaz_length, lengths)
        reverse_batch_gaz_mask = reverse_padded_sequence(reverse_batch_gaz_mask, lengths)

        # word embedding (32,117,50)
        char_embs = self.char_embedding(char_inputs)
        reverse_char_embs = reverse_padded_sequence(char_embs, lengths)

        # gaz embedding (32,117,50)
        gaz_embs = self.gaz_embed((batch_gaz_ids, batch_gaz_length, batch_gaz_mask))
        reverse_gaz_embs = self.gaz_embed((reverse_batch_gaz_ids, reverse_batch_gaz_length, reverse_batch_gaz_mask))
        # print(gaz_embs[0][0][:20])

        # lstm (32,117,50*2)
        forward_inputs = torch.cat((char_embs, gaz_embs), dim=-1)
        backward_inputs = torch.cat((reverse_char_embs, reverse_gaz_embs), dim=-1)

        lstm_outs, _ = self.lstm((forward_inputs, backward_inputs), char_seq_lengths)

        # hidden2tag
        outs = self.hidden2tag(lstm_outs)

        # crf and loss
        loss = self.crf.neg_log_likelihood_loss(outs, mask, batch_label)
        _, tag_seq = self.crf._viterbi_decode(outs, mask)

        return loss, tag_seq

    def forward(self, gaz_list, reverse_gaz_list, char_inputs, char_seq_lengths, mask):
        """
        Args:
            gaz_list: the forward gaz_list
            reverse_gaz_list: the backward gaz list
            char_inputs: word ids
            char_seq_lengths: each sentence length
            mask: sentence mask
        """
        batch_size = char_inputs.size(0)
        seq_len = char_inputs.size(1)
        lengths = list(map(int, char_seq_lengths))

        # get batch gaz ids
        batch_gaz_ids, batch_gaz_length, batch_gaz_mask = get_batch_gaz(reverse_gaz_list, batch_size, seq_len, self.gpu)
        reverse_batch_gaz_ids, reverse_batch_gaz_length, reverse_batch_gaz_mask = get_batch_gaz(gaz_list, batch_size, seq_len, self.gpu)
        reverse_batch_gaz_ids = reverse_padded_sequence(reverse_batch_gaz_ids, lengths)
        reverse_batch_gaz_length = reverse_padded_sequence(reverse_batch_gaz_length, lengths)
        reverse_batch_gaz_mask = reverse_padded_sequence(reverse_batch_gaz_mask, lengths)

        # word embedding
        char_embs = self.char_embedding(char_inputs)
        reverse_char_embs = reverse_padded_sequence(char_embs, lengths)

        # gaz embedding
        gaz_embs = self.gaz_embed((batch_gaz_ids, batch_gaz_length, batch_gaz_mask))
        reverse_gaz_embs = self.gaz_embed((reverse_batch_gaz_ids, reverse_batch_gaz_length, reverse_batch_gaz_mask))

        # lstm
        forward_inputs = torch.cat((char_embs, gaz_embs), dim=-1)
        backward_inputs = torch.cat((reverse_char_embs, reverse_gaz_embs), dim=-1)

        lstm_outs, _ = self.lstm((forward_inputs, backward_inputs), char_seq_lengths)

        # hidden2tag
        outs = self.hidden2tag(lstm_outs)

        # crf and loss
        _, tag_seq = self.crf._viterbi_decode(outs, mask)

        return tag_seq


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from numpy import *
from ray import imio, morpho, evaluate as ev
from scipy.ndimage import binary_dilation, generate_binary_structure
import getopt
import errno
import os
import sys
import json
import argparse

import pdb

larva_gt='/groups/flyem/proj/data/larva_fib_groundtruth_round2.h5'
medulla_gt='/groups/flyem/proj/data/larva_fib_groundtruth_round2.h5'

class Error(Exception):
    """Base-class for exceptions in this module."""

class UsageError(Error):
    def __init__(self, msg):
        self.msg = msg

def main(argv):
    parser = argparse.ArgumentParser(description='Perform VI accuracy comparison between two stacks')
    parser.add_argument('--stack1', type=str, help='Stack compared against base stack (stack path or h5)', dest='stack1', required=True)
    parser.add_argument('--synapsejson', type=str, help='JSON file with synapse annotations', dest='synapsejson', default='')
    parser.add_argument('--grayscale', type=str, help='grayscale h5 volume', dest='grayscale', required=True)
    parser.add_argument('--outdir', type=str, help='raveler output directory', dest='outdir', required=True)

    parser.add_argument('--relabel1', action='store_true', default=False, help='Relabel ids in stack 1 to prevent memory issues', dest='relabel1')
    parser.add_argument('--relabelbase', action='store_true', default=False, help='Relabel ids in stack base to prevent memory issues', dest='relabelbase')
    parser.add_argument('--filtersize', type=int, default=1000, help='Size below which comparisons are not made', dest='filtersize')
    parser.add_argument('--VIthres', type=float, default=0.02, help='VI difference below which is unimportant', dest='VIthres')
    parser.add_argument('--stack1_outjson', type=str, default='stack1_merged.json', help='output json file name for stack 1 (outputs over-merged bodies in stack 1)', dest='stack1_outjson')
    parser.add_argument('--stackbase_outjson', type=str, default='stackbase_merged.json', help='output json file name for stack base (outputs over-merged bodies in stack base)', dest='stackbase_outjson')

    parser.add_argument('--dilate1', type=int, default=0, help='Create and dilate edges on stack 1', dest='dilate1')
    parser.add_argument('--dilatebase', type=int, default=0, help='Create and dilate edges on base image', dest='dilatebase')

    args = parser.parse_args()

    filtersize = args.filtersize
    if args.synapsejson != '':
        filtersize = 0
    #pdb.set_trace()
    stack1=args.stack1
    stack1_int = None
    if stack1.endswith('.h5'):
        stack1_int = imio.read_image_stack(stack1)
    else:
        stack1_int = imio.raveler_to_labeled_volume(stack1, glia=False, use_watershed=False)
    stack1_int = morpho.remove_small_connected_components(stack1_int, min_size=filtersize, in_place=True)

    if args.dilate1 > 0:
        conn = 1 # the connectivity, can be in {1, 2, ..., seg.ndim}
        strel = generate_binary_structure(stack1_int.ndim, conn)
        bound = morpho.seg_to_bdry(stack1_int)
        bound = binary_dilation(bound, strel, iterations=args.dilate1)
        stack1_int[bound] = 0

    if args.relabel1:
        stack1_int, dummy1, dummy2 = ev.relabel_from_one(stack1_int)
    print "Imported first stack with {0} superpixels".format(unique(stack1_int).shape[0])
    #pdb.set_trace()
    gt_fname = args.grayscale
    out_dirname = args.outdir
    gg = imio.read_image_stack(gt_fname,'volume/data')
    if len(gg.shape)==3:
        imio.raveler_output_shortcut(stack1_int+1,stack1_int+1,gg,out_dirname)
    else:
        imio.raveler_output_shortcut(stack1_int+1,stack1_int+1,gg[0,:,:,:],out_dirname)

sys.exit(main(sys.argv))


import torch
import cv2
import numpy as np
from torch.autograd import Variable
from darknet import Darknet
from util import process_result, load_images, resize_image, cv_image2tensor, transform_result
import pickle as pkl
import argparse
import math
import random
import os.path as osp
import os
import sys
from datetime import datetime
from tqdm import tqdm
import ast

def load_classes(namesfile):
    fp = open(namesfile, "r")
    names = fp.read().split("\n")
    return names

def parse_args():
    parser = argparse.ArgumentParser(description='YOLOv3 object detection')
    parser.add_argument('-i', '--input', required=True, help='input image or directory or video')
    parser.add_argument('-t', '--obj-thresh', type=float, default=0.5, help='objectness threshold, DEFAULT: 0.5')
    parser.add_argument('-n', '--nms-thresh', type=float, default=0.4, help='non max suppression threshold, DEFAULT: 0.4')
    parser.add_argument('-o', '--outdir', default='detection', help='output directory, DEFAULT: detection/')
    parser.add_argument('-v', '--video', action='store_true', default=False, help='flag for detecting a video input')
    parser.add_argument('-w', '--webcam', action='store_true',  default=False, help='flag for detecting from webcam. Specify webcam ID in the input. usually 0 for a single webcam connected')
    parser.add_argument('--cuda', action='store_true', default=False, help='flag for running on GPU')
    parser.add_argument('--no-show', action='store_true', default=False, help='do not show the detected video in real time')
    parser.add_argument('-c','--blurclasses', type=list, default=['person','truck', 'car','motorbike'], help='list of objects to blur')
    parser.add_argument('-b','--blurcolor', type=str, default='blur', help='blur or color')
    args = parser.parse_args()

    return args

def create_batches(imgs, batch_size):
    num_batches = math.ceil(len(imgs) // batch_size)
    batches = [imgs[i*batch_size : (i+1)*batch_size] for i in range(num_batches)]

    return batches

def draw_bbox(imgs, bbox, colors, classes,blur_classes,blur_color,read_frames,output_path):
    img = imgs[int(bbox[0])]

    label = classes[int(bbox[-1])]

    if label in blur_classes:

        confidence = int(float(bbox[6])*100)

        label = label+' '+str(confidence)+'%'

        p1 = tuple(bbox[1:3].int())
        p2 = tuple(bbox[3:5].int())

        if blur_color=='blur':
            topLeft = p1
            bottomRight = p2
            x, y = topLeft[0], topLeft[1]
            w, h = bottomRight[0] - topLeft[0], bottomRight[1] - topLeft[1]

            # Grab ROI with Numpy slicing and blur
            ROI = img[y:y+h, x:x+w]
            blur = cv2.GaussianBlur(ROI, (99,99), 0)

            # Insert ROI back into image
            img[y:y+h, x:x+w] = blur
        else:
            blur_color = ast.literal_eval(blur_color)
            cv2.rectangle(img, p1, p2, blur_color, -1)

def detect_video(model, args):

    input_size = [int(model.net_info['height']), int(model.net_info['width'])]

    colors = pkl.load(open("pallete", "rb"))
    classes = load_classes("cfg/coco.names")
    blur_classes = args.blurclasses
    blur_color = args.blurcolor

    if args.webcam:
        cap = cv2.VideoCapture(0)
        output_path = osp.join(args.outdir, 'det_webcam.avi')
    else:
        cap = cv2.VideoCapture(args.input)
        output_path = args.outdir

    width, height = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    fps = cap.get(cv2.CAP_PROP_FPS)

    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    read_frames = 0

    start_time = datetime.now()
    print('Detecting...')
    while cap.isOpened():
        retflag, frame = cap.read()
        read_frames += 1
        if read_frames>0:
            if retflag:
                frame_tensor = cv_image2tensor(frame, input_size).unsqueeze(0)
                frame_tensor = Variable(frame_tensor)

                if args.cuda:
                    frame_tensor = frame_tensor.cuda()

                detections = model(frame_tensor, args.cuda).cpu()
                detections = process_result(detections, args.obj_thresh, args.nms_thresh)
                if len(detections) != 0:
                    detections = transform_result(detections, [frame], input_size)

                    for detection in detections:

                        draw_bbox([frame], detection, colors, classes,blur_classes,blur_color,read_frames,output_path)

                if not args.no_show:
                    cv2.imshow('frame', frame)
                out.write(frame)
                if read_frames % 30 == 0:
                    print('Number of frames processed:', read_frames)
                if not args.no_show and cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            else:
                break

    end_time = datetime.now()
    print('Detection finished in %s' % (end_time - start_time))
    print('Total frames:', read_frames)
    cap.release()
    out.release()
    if not args.no_show:
        cv2.destroyAllWindows()

    print('Detected video saved to ' + output_path)

    return

def detect_image(model, args):

    print('Loading input image(s)...')
    input_size = [int(model.net_info['height']), int(model.net_info['width'])]
    batch_size = int(model.net_info['batch'])

    imlist, imgs = load_images(args.input)
    print('Input image(s) loaded')

    img_batches = create_batches(imgs, batch_size)

    # load colors and classes
    colors = pkl.load(open("pallete", "rb"))
    classes = load_classes("cfg/coco.names")
    blur_classes = args.blurclasses
    blur_color = args.blurcolor

    if not osp.exists(args.outdir):
        os.makedirs(args.outdir)

    start_time = datetime.now()
    print('Detecting...')

    for batchi, img_batch in tqdm(enumerate(img_batches)):
        img_tensors = [cv_image2tensor(img, input_size) for img in img_batch]
        img_tensors = torch.stack(img_tensors)
        img_tensors = Variable(img_tensors)
        if args.cuda:
            img_tensors = img_tensors.cuda()
        detections = model(img_tensors, args.cuda).cpu()
        detections = process_result(detections, args.obj_thresh, args.nms_thresh)
        if len(detections) == 0:
            continue

        detections = transform_result(detections, img_batch, input_size)

        for detection in detections:
            draw_bbox(img_batch, detection, colors, classes,blur_classes,blur_color,0,args.outdir)

        for i, img in enumerate(img_batch):
            save_path = osp.join(args.outdir, osp.basename(imlist[batchi*batch_size + i]))
            cv2.imwrite(save_path, img)
            print(save_path, 'saved')

    end_time = datetime.now()
    print('Detection finished in %s' % (end_time - start_time))

    return

def main():

    args = parse_args()

    if args.cuda and not torch.cuda.is_available():
        print("ERROR: cuda is not available, try running on CPU")
        sys.exit(1)

    print('Loading network...')
    model = Darknet("cfg/yolov3.cfg")
    model.load_weights('weights/yolov3.weights')
    if args.cuda:
        model.cuda()

    model.eval()
    print('Network loaded')

    if args.video:
        detect_video(model, args)

    else:
        detect_image(model, args)

if __name__ == '__main__':
    main()