#!/usr/bin/env pyth
# -*- coding: utf-8 -*-
"""
Created on Thu Feb 18 15:45:32 2021

@author: michael
"""
from datetime import datetime
import torch
import xitorch
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import argparse
import os
from torch import nn, optim
from LoadData import LoadData
import shutil
from _utils import EarlyStopping
from torch.utils.tensorboard import SummaryWriter
from typing import Callable, Union, Mapping, Any, Sequence, Optional
from xitorch._utils.assertfuncs import assert_fcn_params, assert_runtime
from xitorch._core.pure_function import get_pure_function, make_sibling
from xitorch._impls.integrate.ivp.explicit_rk import rk4_ivp, rk38_ivp
from xitorch._impls.integrate.ivp.adaptive_rk import rk23_adaptive, rk45_adaptive
from xitorch._utils.misc import set_default_option, TensorNonTensorSeparator, TensorPacker
from xitorch._utils.tensor import convert_none_grads_to_zeros
from xitorch._docstr.api_docstr import get_methods_docstr
from xitorch.debug.modes import is_debug_enabled

__all__ = ["solve_ivp"]

def solve_ivp(fcn: Union[Callable[..., torch.Tensor], Callable[..., Sequence[torch.Tensor]]],
              ts: torch.Tensor,
              y0: torch.Tensor,
              params: Sequence[Any] = [],
              bck_options: Mapping[str, Any] = {},
              method: Optional[str] = None,
              **fwd_options) -> Union[torch.Tensor, Sequence[torch.Tensor]]:
    r"""
    Solve the initial value problem (IVP) or also commonly known as ordinary
    differential equations (ODE), where given the initial value :math:`\mathbf{y_0}`,
    it then solves

    .. math::

        \mathbf{y}(t) = \mathbf{y_0} + \int_{t_0}^{t} \mathbf{f}(t', \mathbf{y}, \theta)\ \mathrm{d}t'

    Arguments
    ---------
    fcn: callable
        The function that represents dy/dt. The function takes an input of a
        single time ``t`` and tensor ``y`` with shape ``(*ny)`` and
        produce :math:`\mathrm{d}\mathbf{y}/\mathrm{d}t` with shape ``(*ny)``.pip3 insta;;
        The output of the function must be a tensor with shape ``(*ny)`` or
        a list of tensors.
    ts: torch.tensor
        The time points where the value of `y` will be returned.
        It must be monotonically increasing or decreasing.
        It is a tensor with shape ``(nt,)``.
    y0: torch.tensor
        The initial value of ``y``, i.e. ``y(t[0]) == y0``.
        It is a tensor with shape ``(*ny)`` or a list of tensors.
    params: list
        Sequence of other parameters required in the function.
    bck_options: dict
        Options for the backward solve_ivp method. If not specified, it will
        take the same options as fwd_options.
    method: str or None
        Initial value problem solver. If None, it will choose ``"rk45"``.
    **fwd_options
        Method-specific option (see method section below).

    Returns
    -------
    torch.tensor or a list of tensors
        The values of ``y`` for each time step in ``ts``.
        It is a tensor with shape ``(nt,*ny)`` or a list of tensors
    """
    if is_debug_enabled():
        assert_fcn_params(fcn, (ts[0], y0, *params))
    assert_runtime(len(ts.shape) == 1, "Argument ts must be a 1D tensor")

    if method is None:  # set the default method
        method = "rk45"
    fwd_options["method"] = method

    # run once to see if the outputs is a tuple or a single tensor
    is_y0_list = isinstance(y0, list) or isinstance(y0, tuple)
    dydt = fcn(ts[0], y0, *params)
    is_dydt_list = isinstance(dydt, list) or isinstance(dydt, tuple)
    if is_y0_list != is_dydt_list:
        raise RuntimeError("The y0 and output of fcn must both be tuple or a tensor")

    pfcn = get_pure_function(fcn)
    if is_y0_list:
        nt = len(ts)
        roller = TensorPacker(y0)

        @make_sibling(pfcn)
        def pfcn2(t, ytensor, *params):
            ylist = roller.pack(ytensor)
            res_list = pfcn(t, ylist, *params)
            res = roller.flatten(res_list)
            return res

        y0 = roller.flatten(y0)
        res = _SolveIVP.apply(pfcn2, ts, fwd_options, bck_options, len(params), y0, *params, *pfcn.objparams())
        return roller.pack(res)
    else:
        return _SolveIVP.apply(pfcn, ts, fwd_options, bck_options, len(params), y0, *params, *pfcn.objparams())

class _SolveIVP(torch.autograd.Function):
    @staticmethod
    def forward(ctx, pfcn, ts, fwd_options, bck_options, nparams, y0, *allparams):
        config = fwd_options
        ctx.bck_config = set_default_option(config, bck_options)

        params = allparams[:nparams]
        objparams = allparams[nparams:]

        orig_method = config.pop("method")
        method = orig_method.lower()
        try:
            solver = {
                "rk4": rk4_ivp,
                "rk38": rk38_ivp,
                "rk23": rk23_adaptive,
                "rk45": rk45_adaptive,
            }[method]
        except KeyError:
            raise RuntimeError("Unknown solve_ivp method: %s" % config["method"])
        yt = solver(pfcn, ts, y0, params, **config)
        # print(yt)

        #########################
        #####multiprocessing#####
        #########################

        # save the parameters for backward
        ctx.param_sep = TensorNonTensorSeparator(allparams, varonly=True)
        tensor_params = ctx.param_sep.get_tensor_params()
        ctx.save_for_backward(ts, y0, *tensor_params)
        ctx.pfcn = pfcn
        ctx.nparams = nparams
        ctx.yt = yt
        ctx.ts_requires_grad = ts.requires_grad

        return yt

    @staticmethod
    def backward(ctx, grad_yt):
        # grad_yt: (nt, *ny)
        nparams = ctx.nparams
        pfcn = ctx.pfcn
        param_sep = ctx.param_sep
        yt = ctx.yt
        ts_requires_grad = ctx.ts_requires_grad

        # restore the parameters
        saved_tensors = ctx.saved_tensors
        ts = saved_tensors[0]
        y0 = saved_tensors[1]
        tensor_params = list(saved_tensors[2:])
        allparams = param_sep.reconstruct_params(tensor_params)  # maybe change the device
        ntensor_params = len(tensor_params)
        params = allparams[:nparams]
        objparams = allparams[nparams:]

        grad_enabled = torch.is_grad_enabled()

        # custom function to evaluate the input `pfcn` based on whether we want
        # to connect the graph or not
        def pfunc2(t, y, tensor_params):
            if not grad_enabled:
                # if graph is not constructed, then use the default tensor_params
                ycopy = y.detach().requires_grad_()  # [yi.detach().requires_grad_() for yi in y]
                tcopy = t.detach().requires_grad_(False)
                f = pfcn(tcopy, ycopy, *params)
                return f, ycopy, tensor_params
            else:
                # if graph is constructed, then use the clone of the tensor params
                # so that infinite loop of backward can be avoided
                tensor_params_copy = [p.clone().requires_grad_() for p in tensor_params]
                ycopy = y.clone().requires_grad_()
                tcopy = t.clone().requires_grad_(False)
                allparams_copy = param_sep.reconstruct_params(tensor_params_copy)
                params_copy = allparams_copy[:nparams]
                objparams_copy = allparams_copy[nparams:]
                with pfcn.useobjparams(objparams_copy):
                    f = pfcn(tcopy, ycopy, *params_copy)
                return f, ycopy, tensor_params_copy

        # slices and indices definitions on the augmented states
        y_index = 0
        dLdy_index = 1
        # dLdt_index = 2
        # dLdt_slice = slice(dLdt_index, dLdt_index+1, None) # [2:3]
        dLdp_slice = slice(-ntensor_params, None, None) if ntensor_params > 0 else slice(0, 0,
                                                                                None)  # [-ntensor_params:]
        state_size = 2 + ntensor_params  # 3 + ntensor_params
        states = [None for _ in range(state_size)]  # .to(ts.device())

        def new_pfunc(t, states, *tensor_params):
            # t: single-element
            y = states[y_index]
            dLdy = -states[dLdy_index]
            with torch.enable_grad():
                f, y2, tensor_params2 = pfunc2(t, y, tensor_params)
            allgradinputs = ([y2] + list(tensor_params2))
            # allgradinputs = list(tensor_params2)
            allgrads = torch.autograd.grad(f,
                                           inputs=allgradinputs,
                                           grad_outputs=dLdy,
                                           retain_graph=True,
                                           allow_unused=True,
                                           create_graph=torch.is_grad_enabled())  # list of (*ny)
            allgrads = convert_none_grads_to_zeros(allgrads, allgradinputs)
            outs = (
                f,  # dydt
                *allgrads,
            )
            return outs

        ts_flip = ts.flip(0)
        t_flip_idx = -1
        states[y_index] = yt[t_flip_idx]
        states[dLdy_index] = grad_yt[t_flip_idx]
        # states[dLdt_index] = torch.zeros_like(ts[0])
        states[dLdp_slice] = [torch.zeros_like(tp) for tp in tensor_params]
        grad_ts = [None for _ in range(len(ts))] if ts_requires_grad else None

        for i in range(len(ts_flip) - 1):
            # sprint(i)
            # print(states[-1])
            t_flip_idx -= 1
            # ctx.bck_config["methods"] = 'rk4'
            outs = solve_ivp(new_pfunc, ts_flip[i:i + 2], states, tensor_params,
                             fwd_options=ctx.bck_config, bck_options=ctx.bck_config, method="rk4")
            # only take the output for the earliest time
            states = [out[-1] for out in outs]
            states[y_index] = yt[t_flip_idx]
            # gyt is the contribution from the input grad_y
            # gy0 is the propagated gradients from the later time step
            states[dLdy_index] = grad_yt[t_flip_idx] + states[dLdy_index]
        for ii in range(len(states)):
            states[ii] = states[ii].to(ts.device)

        # if ts_requires_grad:
        #     grad_ts[0] = states[dLdt_index].reshape(-1)

        grad_y0 = states[dLdy_index]  # dL/dy0, (*ny)
        if ts_requires_grad:
            grad_ts = torch.cat(grad_ts).reshape(*ts.shape)
        grad_tensor_params = states[dLdp_slice]
        grad_ntensor_params = [None for _ in range(len(allparams) - ntensor_params)]
        grad_params = param_sep.reconstruct_params(grad_tensor_params, grad_ntensor_params)
        return (None, grad_ts, None, None, None, grad_y0, *grad_params)

# docstring completion
ivp_methods = {
    "rk45": rk45_adaptive,
    "rk23": rk23_adaptive,
    "rk4": rk4_ivp,
}
solve_ivp.__doc__ = get_methods_docstr(solve_ivp, ivp_methods)

###################################################
################### GLV model #####################
###################################################

class CancerODEGlv_CPU(xitorch.EditableModule):

    def __init__(self, patientNo, **params):
        # self.r = params["r"]
        self.A = params["A"]
        self.K = params["K"]
        self.pars = params["pars"]
        self.data = LoadData()._Patient_data(patientNo)
        if patientNo == "patient002":
            self.data = self.data[:84]
        if patientNo == "patient046":
            self.data[43:46, 1] -= 10
        if patientNo == "patient056":
            self.data[46, 1] = (self.data[44, 1] + self.data[48, 1]) / 2
        if patientNo == "patient086":
            self.data[1, 1] = (self.data[1, 1] + self.data[8, 1]) / 2
        if patientNo == "patient104":
            self.data = self.data[:(-3)]
        # normalization of drug
        self.CPA = torch.from_numpy(self.data[:, 2]).float()
        self.LEU = torch.from_numpy(self.data[:, 3]).float()
        self.Days = torch.from_numpy(self.data[:, 6] - self.data[0, 6]).float()
        self.OnOff = self.data[:, 5];
        # self.pre_leu()
        self.Response = self.drug_response(
            torch.linspace(start=self.Days[0], end=self.Days[-1], steps=int(self.Days[-1] - self.Days[0]) + 1))
        self.cell_size = 5.236e-10  # 4. / 3. * 3.1415926 * (5e-4cm) ** 3   # cm^3

    def drug_response(self, t):
        drug = torch.zeros((t.shape[0], 2), dtype=torch.float)
        slice0 = torch.bucketize(t, self.Days, right=True) - 1
        slice_75 = torch.where(self.LEU == 7.5)[0]
        dose75_date = self.Days[slice_75]
        slice_225 = torch.where(self.LEU == 22.5)[0]
        dose225_date = self.Days[slice_225]
        slice_30 = torch.where(self.LEU == 30)[0]
        dose30_date = self.Days[slice_30]
        slice_45 = torch.where(self.LEU == 45)[0]
        dose45_date = self.Days[slice_45]

        I0_CPA_dose = self.CPA[slice0]
        I0_LEU_dose = self.LEU[slice0] # LEu dose has 7.5/22.5/30/45, 4 different dosages, for 4/12/16/24 weeks, and we can see that no matter
        dose75 = dose225 = dose30 = dose45 = 1
        _date = -100
        for date in dose75_date.int():
            if abs(date - _date) < 7 or abs(date - _date) == 7:
                dose75 += 1
            else:
                dose75 = 1
            if dose75 == 1:
                temp = torch.zeros(28, dtype = torch.float)  # last 12 weeks
                temp[0:7 * 1] = - 3.75 / 6 * torch.arange(0, 7, 1)
                temp[(7 * 1):(7 * 3)] = (7.5 + 3.75) / (20 - 6) * torch.arange(7, 21, 1) + (
                            7.5 - (7.5 + 3.75) / (20 - 6) * 20)
                temp[(7 * 3):] = 7.5
                I0_LEU_dose[date: (date + 7 * 4)] = temp[0:I0_LEU_dose[date: (date + 7 * 4)].size(0)]
            else:
                I0_LEU_dose[_date: (date + 7 * 4)] = 7.5
            _date = date + 7 * 4
        _date = -100
        for date in dose225_date.int():
            if abs(date - _date) <  7 or abs(date - _date) == 7:
                dose225 += 1
            else:
                dose225 = 1
            if dose225 == 1:
                temp = torch.zeros(7*12, dtype = torch.float) # last 12 weeks
                temp[0:7*1] = - 3.75/6 * torch.arange(0,7,1)
                temp[(7*1):(7*3)] = (7.5 + 3.75)/(20-6) *torch.arange(7, 21, 1) +(7.5 -  (7.5 + 3.75)/(20-6)*20)
                temp[(7*3):] = 7.5
                I0_LEU_dose[date: (date + 7 * 12)] = temp[0:I0_LEU_dose[date: (date + 7 * 12)].size(0)]
            else:
                I0_LEU_dose[_date: (date + 7 * 12)] = 7.5
            _date = date + 7 * 12

        _date = -100
        for date in dose30_date.int():
            if abs(date - _date) <  7 or abs(date - _date) == 7:
                dose30 += 1
            else:
                dose30 = 1
            if dose30 == 1:
                temp = torch.zeros(7*16, dtype = torch.float)  # last 12 weeks
                temp[0:7 * 1] = - 3.75 / 6 * torch.arange(0, 7, 1)
                temp[(7 * 1):(7 * 3)] = (7.5 + 3.75) / (20 - 6) * torch.arange(7, 21, 1) + (
                            7.5 - (7.5 + 3.75) / (20 - 6) * 20)
                temp[(7 * 3):] = 7.5
                I0_LEU_dose[date: (date + 7 * 16)] = temp[0:I0_LEU_dose[date: (date + 7 * 16)].size(0)]
            else:
                I0_LEU_dose[_date: (date + 7 * 16)] = 7.5
            _date = date + 7 * 16
        _date = -100
        for date in dose45_date.int():
            if abs(date - _date) <  7 or abs(date - _date) == 7:
                dose45 += 1
            else:
                dose45 = 1
            if dose45 == 1:
                temp = torch.zeros(7 * 24, dtype = torch.float)  # last 12 weeks
                temp[0:7 * 1] = - 3.75 / 6 * torch.arange(0, 7, 1)
                temp[(7 * 1):(7 * 3)] = (7.5 + 3.75) / (20 - 6) * torch.arange(7, 21, 1) + (
                            7.5 - (7.5 + 3.75) / (20 - 6) * 20)
                temp[(7 * 3):] = 7.5
                I0_LEU_dose[date: (date + 7 * 24)] = temp[0:I0_LEU_dose[date: (date + 7 * 24)].size(0)]
            else:
                I0_LEU_dose[_date: (date + 7 * 24)] = 7.5
            _date = date + 7 * 24

        I0_LEU_dose[torch.cat((torch.where(I0_LEU_dose == 22.5)[0],torch.where(I0_LEU_dose == 30)[0],torch.where(I0_LEU_dose == 45)[0]))] = 7.5
        I0_LEU_dose = I0_LEU_dose / 7.5
        I0_CPA_dose = I0_CPA_dose / 200
        for ii in range(1, I0_LEU_dose.shape[0] - 1):
            if  I0_LEU_dose[ii+1] < 0 and I0_LEU_dose[ii-1] == 1.:
                I0_LEU_dose[ii : (ii + 7 * 3)] = 1.
        drug[:, 0] = I0_CPA_dose
        drug[:, 1] = I0_LEU_dose
        return drug

    def forward(self, t, y):
        r = self.pars[0:2]
        beta = self.pars[2:4]
        Beta = torch.zeros((2, 2), dtype=torch.float)
        Beta[:, 0] = beta
        phi = self.pars[-4]
        betac = self.pars[(-2):]
        A = torch.tensor([1., .5, .5, 1.]).view(2, 2)
        A[0,1] = 1/(1 + torch.exp(-self.pars[-3] * torch.tensor([t /28/12]))) #+= self.pars[-3] * t /28/12 #
        # A[0, 1] = a1; A[1, 0] = a2
        gamma = 0.25  # the half life time for psa is 2.5 days, so each day psa decrease by  25%
        # drug = torch.zeros((t.shape[0], 2), dtype = torch.float)
        # slice0 = torch.bucketize(t, self.Days, right = True) - 1
        # I_CPA_dose = self.CPA[slice0];I_LEU_dose = self.LEU[slice0]
        index = int(t) if len(t.shape) == 0 else t.int().cpu().numpy()
        x = y[0:2]  # cell count
        p = y[-1]  # psa level
        # comp = torch.tensor([x[0]*2 + x[1]*a, x[0]*a + x[1]*1])
        dxdt = torch.multiply(
            r * x, (1 - (x @ A / self.K) ** phi - self.Response[index] @ Beta))  # -
        # self.Response[index, 0] * self.Response[index, 1] * Mu))
        dpdt = betac @ x * self.cell_size - gamma * p  # /(abs(sum(dxdt)) + gamma) /(abs(sum(dxdt)*1000/self.K) + gamma)
        df = torch.zeros(3, dtype=torch.float)
        df[0:2], df[-1] = dxdt, dpdt
        return df

    def getparamnames(self, methodname, prefix=""):
        if methodname == "forward":
            return [prefix + "A",  prefix + "K", prefix + "pars"]  # , prefix+"beta", prefix+"gamma", prefix+"alpha"]
        else:
            raise KeyError()

def MSEloss_weight(inputs, targets, weight):
    return torch.sum(weight * (inputs - targets) ** 2)

def clip_grad(grad, max_norm: float, norm_type: float = 2.0) -> torch.Tensor:

    max_norm = float(max_norm)
    norm_type = float(norm_type)
    if len(grad) == 0:
        return torch.tensor(0.)
    device = grad.device
    total_norm = torch.norm(grad.detach(), norm_type).to(device)
    clip_coef = max_norm / (total_norm + 1e-6)
    if clip_coef < 1:
        grad.detach().mul_(clip_coef.to(grad.device))
    return total_norm

# As for the determination of the initial value, we define if psa =10, the cell count for AD is 1e8

from collections import deque

def train_glv(args, alldata):

    i = args.number
    fail_deque =  deque(maxlen = 10)
    alpha = 0.25 #Alpha[i]
    # under_i = [4, 13, 19, 54, 60, 83, 87, 95, 96, 100, 101, 105]
    # if i in under_i:
    #     alpha = 0.5
    cell_size = 5.236e-10
    fail_flag = False
    if len(str(i)) == 1:
        patientNo = "patient00" + str(i)
    elif len(str(i)) == 2:
        patientNo = "patient0" + str(i)
    else:
        patientNo = "patient" + str(i)
    print(patientNo)

    # PARS = np.array(pd.read_csv("PARS.csv"), dtype=np.float).reshape(-1)
    # if patientNo in ["patient012","patient015" ,"patient006"]:
    #     continue
    data = alldata[patientNo]
    if patientNo == "patient002":
        data = data[:84]
    if patientNo == "patient046":
        data[43:46, 1] -= 10
    if patientNo == "patient056":
        data[46, 1] = (data[44,1] + data[48,1])/2
    if patientNo == "patient086":
        data[1,1] = (data[1,1]+data[8,1])/2
    if patientNo == "patient104":
        data = data[:(-3)]
    Days = data[:, 6] - data[0, 6]
    OnOff = data[:, 5]
    Cycle = data[:,-3]
    PSA = data[:, 1]
    index = np.where(np.isnan(PSA))[0]
    PSA = torch.from_numpy(np.delete(PSA, index)).float()
    DAYS = np.delete(Days, index)
    treatInt = [0.]
    validate_set = deque(maxlen=int(Cycle[-1]))
    for ii in range(1, OnOff.shape[0] - 1):
        if OnOff[ii - 1] == 1 and OnOff[ii] == 0:
            treatInt.append(Days[ii])
        if OnOff[ii - 1] == 0 and OnOff[ii] == 1:
            treatInt.append(Days[ii])
    treatInt.append(Days[-1])
    slicing = np.digitize(treatInt, DAYS, right=True)
    for kk in np.arange(slicing.shape[0]-1, step = 2):
        if kk + 2 < slicing.shape[0]:
            loo = np.random.choice(DAYS[slicing[kk] + 1:slicing[kk + 2]])
        elif kk + 2 >= slicing.shape[0] and DAYS[slicing[kk]+1:].size != 0:
            loo = np.random.choice(DAYS[slicing[kk]+1:])
        else:
            break
        validate_set.append(loo)
    validate_days = np.array(validate_set, dtype=np.int32)
    validate_psa = PSA[np.isin(DAYS, validate_set)].detach().numpy()
    train_days = DAYS[~np.isin(DAYS, validate_set)]
    train_slice = np.digitize(treatInt, train_days, right=True)
    train_psa = PSA[~np.isin(DAYS, validate_set)]
    # 5e+8 5cm^3 = 5000mm^3
    mean_v = 5
    mean_psa = 22.1
    K1 = 1.1 * mean_v * (max(PSA)/mean_psa)/cell_size  # 2e+11
    K2 = alpha * K1
    K = torch.tensor([K1, K2]) #1.1 * mean_v * (max(PSA)/mean_psa)/cell_size #
    # r = torch.tensor([1.2, 1.5], dtype = torch.float) # AD is 2 times more responsive compared to AI cell
    # e = torch.tensor([[0.8], [0.9]], dtype = torch.float) # AD is 5 times more competitive effects exert than AI cells
    A = torch.tensor([1., .5, 0.5, 1.], dtype=torch.float).view(2, 2)
    inputs = torch.linspace(start=Days[0], end=Days[-1], steps=int(Days[-1] - Days[0]) + 1, dtype=torch.float)
    # criterion = nn.MSELoss()
    # learning_rate = torch.tensor([.00001, .00001, .001, .001, .0001, .00001, .001, .001])
    # momentum = 0.98
    # velocity = 0.
    # pars = torch.tensor([0.05, 0.05, 2, 1, 1., -0.01, 1., 1.]).float().requires_grad_()
    parsdir = "./retrain-sigmoid/model_pars"
    parslist = os.listdir(parsdir)
    PARS = torch.zeros(8, dtype = torch.float)
    # reading the ode parameters and the initial/terminal states
    for file in parslist:
        pars_df = pd.read_csv(parsdir + '/' + file)
        pars = torch.from_numpy(np.array(pars_df.loc[4, ~np.isnan(pars_df.loc[4, :])])).float()
        PARS += pars

    pars = PARS/len(parslist)
    pars[-3] = 0
    # if patientNo in ["patient063", "patient075", "patient083", "patient050", "patient036", "patient088", "patient105", "patient087", "patient095", "patient101",
    #                  "patient001", "patient003", "patient004", "patient011", "patient013", "patient046", "patient063", "patient036", "patient062", "patient046",
    #                  "patient013", "patient016", "patient017", "patient056", "patient071", "patient078", "patient079", "patient091", "patient104", "patient105",
    #                  "patient032", "patient083", "patient092"]:
    #     pars_df = pd.read_csv("./retrain-sigmoid/model_pars" + '/Args_' + patientNo + ".csv")
    #     pars = torch.from_numpy(np.array(pars_df.loc[4, ~np.isnan(pars_df.loc[4, :])])).float()
    pars = pars.requires_grad_()
    inits_pars = pars.detach().numpy()
    loss_deque = deque(maxlen=10)
    loss = torch.tensor([10000], dtype=torch.float)
    Epoch = 10000
    best_loss = 10000
    best_pars = pars.detach().numpy()
    cancerode = CancerODEGlv_CPU(patientNo, A=A, K=K, pars=pars)
    optimizer = torch.optim.Adam([cancerode.pars], lr = .001)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size =1000, gamma=0.8)
    # initialize the early_stopping object
    early_stopping = EarlyStopping(patience=15)
    if not os.path.exists("./analysis-sigmoid"):
        os.mkdir("./analysis-sigmoid")
    if not os.path.exists("./analysis-sigmoid/model_infos/" + patientNo):
        os.makedirs("./analysis-sigmoid/model_infos/" + patientNo)
    if not os.path.exists("./analysis-sigmoid/model_plots/" + patientNo):
        # shutil.rmtree("./retrain-sigmoid/model_plots/" + patientNo)
        os.makedirs("./analysis-sigmoid/model_plots/" + patientNo)
    if not os.path.exists("./analysis-sigmoid/model_pars/" + patientNo):
        os.makedirs("./analysis-sigmoid/model_pars/" + patientNo)
    if not os.path.exists("./analysis-sigmoid/model_validate"):
        os.makedirs("./analysis-sigmoid/model_validate")
    log = []
    t = datetime.now().strftime("%Y%m%d-%H%M")
    summary_dir = "./analysis-sigmoid/model_infos" + '/' + str(patientNo) + '/' + str(t) + "/" + str(args.t)
    writer = SummaryWriter(log_dir=summary_dir)
    for epoch in range(Epoch):
        Init = torch.tensor([mean_v/mean_psa * PSA[0]/cell_size, 1e-4 * K2, PSA[0]], dtype=torch.float)
        _loss = loss.detach().numpy()
        _pars = pars.detach().numpy()
        loss = torch.zeros(1, dtype=torch.float)

        res = Init.detach().numpy().reshape(1, -1)
        optimizer.zero_grad()
        for ii in range(len(treatInt) - 1):
            if ii == len(treatInt) - 2:
                INPUTS = inputs[int(treatInt[ii]):]
            else:
                INPUTS = inputs[int(treatInt[ii]):int(treatInt[ii + 1])]
            ts = INPUTS.requires_grad_(False)
            OUT = solve_ivp(cancerode.forward, ts=ts, y0=Init, params=(), atol=1e-08, rtol=1e-05)
            res = np.append(res, OUT.detach().numpy(), axis=0)
            d = train_days[train_slice[ii]:train_slice[ii + 1]]
            slicing1 = np.digitize(d, INPUTS, right=True)
            EST_PSA = OUT[slicing1, -1]
            EST_PSA[torch.isnan(EST_PSA)] = 1000
            psa = train_psa[train_slice[ii]: train_slice[ii + 1]]
            if ii == 0:
                psa[0] = Init[-1]
            Init = OUT[-1]
            weights = 1.
            n_psa = psa.shape[0]
            if ii % 2 == 0:
                weights = torch.linspace(start=1, end=n_psa, steps=n_psa)
            if ii % 2 == 1:
                weights = torch.linspace(start=n_psa, end=1, steps=n_psa)
            weights = weights/sum(weights)
            loss = MSEloss_weight(EST_PSA, psa, weights) #+ loss
            loss.backward(retain_graph=True if ii != len(treatInt) - 2 else False)

        # pars_grad = torch.autograd.grad(loss, (cancerode.pars,))[0]
        # clip_grad(pars_grad, max_norm=10);
        # velocity = momentum * velocity - learning_rate * pars_grad # pars_grad
        # pars = pars.add(velocity)  # .detach().requires_grad_()
        # optimizer.zero_grad()
        # loss.backward()
        torch.nn.utils.clip_grad_norm_(pars, max_norm = 100 * (1.001 - epoch/Epoch))
        optimizer.step()
        with torch.no_grad():
            pars[:2].clamp_( min = 5e-3, max = 1e-1)
            pars[2:4].clamp_( min=0)
            pars[-2:].clamp_( min=1e-3)
        scheduler.step()
        # cancerode.pars = pars
        pars_grad = pars.grad
        flag_pars = pars_grad.detach().numpy()[-3]
        # avoiding over-fitted
        loss_deque.append(_loss)
        print(pars.detach().numpy())
        loss_array = np.array(loss_deque, dtype=np.float).reshape(-1)
        if epoch > 10:
            loss_decay_mask = sum((loss_array[1:] - loss_array[:-1]) / loss_array[:-1] > 0)

        if loss.detach().numpy().item() < best_loss:
            best_loss = loss.detach().numpy().item()
            best_pars = pars.detach().numpy()

        with torch.no_grad():
            res = res[1:]
            ad = res[:, 0]
            ai = res[:, 1]
            p = res[:, 2]
            val_psa = p[validate_days]
            validate_loss = np.mean((val_psa - validate_psa) ** 2)
            # log.append([epoch, loss.detach().numpy().item(), validate_loss])
            writer.add_scalar('Loss', loss.detach().numpy().item(), epoch)
            writer.add_scalar('V-Loss', validate_loss, epoch)
            # if epoch % 50 == 0:
            #     file_writing_obj = open('./analysis-sigmoid/model_infos/infos-' + patientNo + "-" + str(args.t) + '.txt', 'w')
            #     file_writing_obj.write(str(log))
            #     file_writing_obj.close()
                #print('Epoch: {} \t Loss: {:.2f} \t Val_Loss: {:.2f}'.format(epoch, loss.detach().numpy().item(), validate_loss))
            if epoch % 100 == 0:
                # x = inputs.detach().numpy()
                # plt.scatter(DAYS, PSA, color="black", marker="*", alpha=0.6)
                # plt.plot(x, p, color="black", linestyle="-", linewidth=1)
                # plt.xlabel("Time (Days)" )
                # plt.ylabel("PSA level (ug/ml)")
                # plt.savefig("./analysis-sigmoid/model_plots/"+patientNo+"/PSA_" + str(args.t) + "-" + patientNo + "_" + str(epoch) + ".png", dpi=100)
                # # plt.show()
                # plt.close()
                # plt.plot(x, ad, color="black", linestyle="--", linewidth=1, label="AD")
                # plt.plot(x, ai, color="black", linestyle="-.", linewidth=1, label="AI")
                # plt.xlabel("Time (Days)")
                # plt.ylabel("Cell counts")
                # plt.legend(loc='upper right')
                # plt.savefig("./analysis-sigmoid/model_plots/"+patientNo+"/Cell_All_" + str(args.t) + "-" + patientNo + "_" + str(epoch) + ".png", dpi=100)
                # # plt.show()
                # plt.close()
                # plt.plot(x, ai, color="black", linestyle="-.", linewidth=1, label="AI")
                # plt.xlabel("Time (Days)")
                # plt.ylabel("Cell counts")
                # plt.savefig("./analysis-sigmoid/model_plots/"+patientNo+"/Cell_AI_" + str(args.t) + "-" + patientNo + "_" + str(epoch) + ".png", dpi=100)
                # # plt.show()
                # plt.close()
                A1 = A.detach().numpy().reshape(-1)
                K1 = K.detach().numpy()
                terminate = res[-1]
                init = Init.detach().numpy()
                states1 = np.append(init, terminate)
                pars_detach = pars.detach().numpy()
                plist = [A1, K1, states1, pars_detach, best_pars]
                plist_df = pd.DataFrame(plist)
                plist_df.to_csv("./analysis-sigmoid/model_pars/"+patientNo+"/Args_" + str(args.t) + "-" + patientNo + ".csv",
                                index=False)
        if epoch > 2000:
            early_stopping(validate_loss, pars)
        if early_stopping.early_stop:
            print("Early stopping")
            break

    Init = torch.tensor([mean_v/mean_psa * PSA[0]/cell_size, 1e-4 * K2, PSA[0]], dtype=torch.float)
    cancerode = CancerODEGlv_CPU(patientNo, A=A, K=K, pars=pars)
    out = solve_ivp(cancerode.forward, ts=inputs, y0=Init, params=(), atol=1e-08, rtol=1e-05)
    ad = out[:, 0].detach().numpy()
    ai = out[:, 1].detach().numpy()
    psa = out[:, -1].detach().numpy()
    pred_validate_psa = psa[validate_days]
    validate_loss = np.array([sum((pred_validate_psa - validate_psa)**2)], dtype = np.float)
    validate_list = [validate_psa, pred_validate_psa, validate_loss]
    validate_df = pd.DataFrame(validate_list, index = ["true", 'predict', 'loss'])
    validate_df.to_csv("./analysis-sigmoid/model_validate/" + patientNo + "/validate_" + str(args.t) + "-" + patientNo + ".csv", index=True)
    print(validate_loss)
    x = inputs.numpy()
    plt.scatter(DAYS, PSA, color="black", marker="*", alpha=0.6)

    plt.plot(x, psa, color="black", linestyle="-", linewidth=1)
    plt.xlabel("Time (Days)")
    plt.ylabel("PSA level (ug/ml)")
    plt.savefig("./analysis-sigmoid/model_plots/"+patientNo+"/Final_PSA_" + str(args.t) + "-" + patientNo + ".png", dpi=300)
    plt.close()

    plt.plot(x, ad, color="black", linestyle="--", linewidth=1, label="AD")
    plt.plot(x, ai, color="black", linestyle="-.", linewidth=1, label="AI")
    plt.xlabel("Time (Days)")
    plt.ylabel("Cell counts")
    plt.legend(loc='upper right')
    plt.savefig("./analysis-sigmoid/model_plots/"+patientNo+"/Final_Cell_" + str(args.t) + "-" + patientNo + ".png", dpi=300)
    plt.close()

    # competition strength
    c = (out[:, :2] @ A[:, 1] / K[1]) ** pars[-3]
    cc = c.detach().numpy()
    plt.plot(x, cc, color="b", linestyle="-", linewidth=1, label="Competition for AI")
    plt.savefig("./analysis-sigmoid/model_plots/"+patientNo+"/Final_Competition_" + str(args.t) + "-" + patientNo + ".png", dpi=300)
    plt.close()

    A = A.detach().numpy().reshape(-1)
    K = K.detach().numpy()
    terminate = out[-1].detach().numpy()
    init = Init.detach().numpy()
    states = np.append(init, terminate)
    pars_detach = pars.detach().numpy()
    plist = [A, K, states, pars_detach, best_pars]
    plist_df = pd.DataFrame(plist)
    plist_df.to_csv("./analysis-sigmoid/model_pars/"+patientNo+"/Args_" + str(args.t) + "-" + patientNo + ".csv", index=False)

parser = argparse.ArgumentParser(description='Patient arguments')
parser.add_argument('--number', '-n', help='Patient No., int type, requested', default=13, type = int)
parser.add_argument('--t',  default=0, type = int)
args = parser.parse_args()

if __name__ == "__main__":
    alldata =  LoadData().Double_Drug()
    train_glv(args, alldata)


from collections  import namedtuple
from itertools    import count
from glob         import glob
from time         import time, sleep

from .synthesis    import Synthesis
from .utils        import Fold
from .kmerSetDB    import kmerSetDB
from .kmerSetArray import kmerSetArray

import sys
import numpy
import uuid
import math
import textwrap

from . import utils
from . import makerchecks
from . import projector
from . import berno
from . import finder

from Bio.SeqUtils import MeltingTemp

class NRPMaker(object):

    def __init__(self, part_type='RNA', seed=None):

        # Instance variables
        self.part_type  = part_type
        self.synthesis  = Synthesis()
        self.fold       = Fold(part_type=part_type)
        self.proj_id    = str(uuid.uuid4())
        self.kmer_db    = None
        self.background = None

        # Seed the RNG
        if not seed is None and isinstance(seed, int):
            self.rng = numpy.random.default_rng(seed=seed)
        else:
            self.rng = numpy.random.default_rng()

        # Lookup tables
        if self.part_type == 'RNA':
            self.iupac_space = {
                'A': {'A'},
                'C': {'C'},
                'G': {'G'},
                'U': {'U'},
                'R': {'A', 'G'},
                'Y': {'C', 'U'},
                'S': {'G', 'C'},
                'W': {'A', 'U'},
                'K': {'G', 'U'},
                'M': {'A', 'C'},
                'B': {'C', 'G', 'U'},
                'V': {'A', 'C', 'G'},
                'D': {'A', 'G', 'U'},
                'H': {'A', 'C', 'U'},
                'N': {'A', 'U', 'G', 'C'}
            }
            self.iupac_compl = {
                'A': 'U', # A - U
                'C': 'G', # C - G
                'G': 'Y', # G - C U - Y
                'U': 'R', # U - A G - R
                'R': 'Y', # R - A G - U C - Y
                'Y': 'R', # Y - C U - G A G - R
                'S': 'B', # S - G C - C G U - B
                'W': 'D', # W - A U - U A G - D
                'K': 'N', # K - G U - C U G A
                'M': 'K', # M - A C - U G - K
                'B': 'N', # B - C G U - G C U A - N
                'V': 'D', # V - A C G - U G A - D
                'D': 'N', # D - A G U - U C A G - N
                'H': 'N', # H - A C U - U G A G - N
                'N': 'N'
            }
            self.base_compl = {
                'A': ['U'],
                'G': ['C', 'U'],
                'C': ['G'],
                'U': ['A', 'G']
            }
        else:
            self.iupac_space = {
                'A': {'A'},
                'C': {'C'},
                'G': {'G'},
                'T': {'T'},
                'R': {'A', 'G'},
                'Y': {'C', 'T'},
                'S': {'G', 'C'},
                'W': {'A', 'T'},
                'K': {'G', 'T'},
                'M': {'A', 'C'},
                'B': {'C', 'G', 'T'},
                'V': {'A', 'C', 'G'},
                'D': {'A', 'G', 'T'},
                'H': {'A', 'C', 'T'},
                'N': {'A', 'T', 'G', 'C'}
            }
            self.iupac_compl = {
                'A': 'T', # A - T
                'C': 'G', # C - G
                'G': 'C', # G - C
                'T': 'A', # T - A
                'R': 'Y', # R - A G - T C - Y
                'Y': 'R', # Y - C T - G A - R
                'S': 'S', # S - G C - C G - S
                'W': 'W', # W - A T - T A - W
                'K': 'M', # K - G T - C A - M
                'M': 'K', # M - A C - T G - K
                'B': 'V', # B - C G T - G C A - V
                'V': 'B', # V - A C G - T G C - B
                'D': 'H', # D - A G T - T C A - H
                'H': 'D', # H - A C T - T G A - D
                'N': 'N'
            }
            self.base_compl = {
                'A': ['T'],
                'G': ['C'],
                'C': ['G'],
                'T': ['A']
            }

    def _get_adjusted_struct(self, struct, seq):
        if struct is None:
            return '.'*len(seq)
        elif len(seq) < len(struct):
            return struct[:len(seq)]
        elif len(seq) > len(struct):
            return ''.join([struct, '.'*(len(seq) - len(struct))])
        return struct

    def _get_meta_struct(self, struct):
        struct_spec   = namedtuple(
            'struct_spec', 'struct paired_dict rev_paired_dict unpaired_set folding inversefolding')
        pairing_stack = []
        meta_struct   = struct_spec(
            struct=struct.replace('x', '.'),
            paired_dict={},
            rev_paired_dict={},
            unpaired_set=set(),
            folding={},
            inversefolding={})

        for index, nt in enumerate(struct):

            if nt == '(':
                pairing_stack.append(index)

            elif nt == ')':

                try:
                    closure = pairing_stack.pop()
                except:
                    raise ValueError(
                        ' [X] Unbalanced parentheses in structure constraint')

                meta_struct.paired_dict[closure]   = index
                meta_struct.rev_paired_dict[index] = closure

            elif nt == 'x':
                meta_struct.unpaired_set.add(index)

        if pairing_stack:
            raise ValueError(
                ' [X] Unbalanced parentheses in structure constraint')

        if meta_struct.paired_dict:
            meta_struct.folding['status'] = True
            meta_struct.inversefolding['status'] = True
        elif meta_struct.unpaired_set:
            meta_struct.folding['status'] = True
            meta_struct.inversefolding['status'] = False
        else:
            meta_struct.folding['status'] = False
            meta_struct.inversefolding['status'] = False

        return meta_struct

    def _get_meta_seq(self, seq, meta_struct):
        seq = list(seq)
        meta_seq = [None] * len(seq)

        # Normalize Meta Sequence
        for i in range(len(seq)):
            try:
                if i in meta_struct.paired_dict:
                    j = meta_struct.paired_dict[i]
                    if len(self.iupac_space[seq[j]]) < len(self.iupac_space[seq[i]]):
                        seq[i] = self.iupac_compl[seq[j]]
                    elif len(self.iupac_space[seq[i]]) < len(self.iupac_space[seq[j]]):
                        seq[j] = self.iupac_compl[seq[i]]
                meta_seq[i] = set(self.iupac_space[seq[i]])
            except:
                raise ValueError(
                    ' [X] Invalid IUPAC code at index {} in sequence constraint'.format(i))

        return tuple(meta_seq)

    def _reset_candidate_kmer_set(
        self,
        candidate,
        kmer_set,
        i):
        candidate[i] = '-'
        kmer_set[i]  = ' '

    def _clear_path(
        self,
        candidate,
        tried_set,
        kmer_set,
        i,
        k):
        j = i
        while j > k:
            self._reset_candidate_kmer_set(
                candidate,
                kmer_set,
                j)
            tried_set[j] = set()
            # tried_set[j] = set(candidate[j])
            j -= 1
        else:
            self._reset_candidate_kmer_set(
                candidate,
                kmer_set,
                j)
        i = j
        return i

    def _get_rollback_index(
        self,
        meta_seq,
        meta_struct,
        i,
        homology,
        candidate,
        tried_set):
        roll_back_index = None
        # See if any position in last homology places has potential for change
        j = i
        while (j >= i - homology + 1) and (j >= 0):
            # Case ( or . with potential for change
            if not (j in meta_struct.rev_paired_dict):
                if len(meta_seq[j]) > len(tried_set[j]):
                    roll_back_index = j
                    break
            # Case ) for RNA Parts with potential for change
            else:
                l = meta_struct.rev_paired_dict[j]
                pnt = meta_seq[j]
                pnt = pnt.intersection(
                    self.base_compl[candidate[l]])
                if len(tried_set[j]) < len(pnt):
                    roll_back_index = j
                    break
            j -= 1
        # Else go to the ( for the first ) in the last homology places
        if not roll_back_index:
            j = max(0, i - homology + 1)
            k = None
            while j <= i:
                if j in meta_struct.rev_paired_dict:
                    k = j
                    break
                j += 1
            roll_back_index = meta_struct.rev_paired_dict[k]
        return roll_back_index

    def _roll_back(
        self,
        meta_seq,
        meta_struct,
        i,
        homology,
        candidate,
        tried_set,
        kmer_set,
        rbi=None):
        # Default Case: Explicit roll back
        if not rbi is None:

            # traceback to ( if corresponding ) given
            if rbi in meta_struct.rev_paired_dict:
                j = meta_struct.rev_paired_dict[rbi]
                pnt = meta_seq[rbi]
                pnt = pnt.intersection(
                    self.base_compl[candidate[j]])
                # No potential for change
                if len(pnt) == len(tried_set[i]):
                    rbi = j

            self._clear_path(
                candidate, tried_set, kmer_set, i, k=rbi)
            return rbi

        # Case 1: i is not paired upstream
        if not i in meta_struct.rev_paired_dict:
            self._reset_candidate_kmer_set(
                candidate, kmer_set, i)

        # Case 2: i is paired upstream
        else:
            roll_back_index = self._get_rollback_index(
                meta_seq, meta_struct, i, homology, candidate, tried_set)
            i = self._clear_path(
                candidate, tried_set, kmer_set, i, k=roll_back_index)

        return i

    def _get_local_roll_back_index(
        self,
        candidate,
        i,
        local_model_fn,
        verbose):
        # Prep candidate
        candidate_str = ''.join(candidate)
        candidate_str = candidate_str[:i+1]

        # Try to evaluate the local_model_fn on candidate_str
        outcome = True
        try:
            outcome = local_model_fn(candidate_str)
            if outcome in [True, False]:
                state, index = outcome, i
            else:
                state, index = outcome
        except Exception as e:
            print(' Local Model fn. failed to evaluate partial path: {}\n'.format(
                candidate_str))
            raise e # No intelligence, halt everything!

        # State satisfactory?
        try:
            assert state in [True, False]
        except Exception as e:
            print(' Local Model fn. failed to evaluate partial path: {}'.format(
                candidate_str))
            print(' Local Model fn. returned a non-boolean evaluation: {}\n'.format(
                state))
            raise e

        # Index satisfactory?
        try:
            if state == False:
                index = int(index)
                assert 0 <= index <= i
        except Exception as e:
            print(' Local Model fn. failed to evaluate partial path: {}'.format(
                candidate_str))
            print(' Local Model fn. returned a non-integer or invalid traceback index: {}\n'.format(
                index))
            raise e

        # No conflict found!
        if state:
            return None
        # Conflict Found!
        else:
            return index

    def _get_non_coding_candidate(
        self,
        meta_seq,
        meta_struct,
        homology,
        local_model_fn,
        verbose,
        jump=False,
        start_seq=None,
        allow_internal_repeat=False):

        # Setup the data structures
        candidate = ['-'] * len(meta_seq) if not start_seq else list(start_seq)
        tried_set = [set() for _ in range(len(meta_seq))]
        kmer_set  = kmerSetArray(size=len(meta_seq))

        # Setup indexing
        i               = 0
        roll_back_count = 0

        # Main backtracking code
        while -1 < i < len(meta_seq):

            # Jumping out of iteration
            if jump and roll_back_count == homology:
                candidate = None
                break

            # Try to build a candidate
            if candidate[i] == '-':

                # Phase determination
                forward = False
                # Case )
                if i in meta_struct.rev_paired_dict:
                    j = meta_struct.rev_paired_dict[i]
                    pnt = meta_seq[i]
                    pnt = pnt.intersection(
                        self.base_compl[candidate[j]])
                    if len(pnt) > len(tried_set[i]):
                        forward = True
                # Case ( and .
                else:
                    if len(tried_set[i]) < len(meta_seq[i]):
                        forward = True

                # Forward phase - A nucleotide may be chosen
                if forward:
                    # Reset roll_back_count
                    roll_back_count = 0

                    # Case ( and .
                    if not i in meta_struct.rev_paired_dict:
                        candidate[i] = self.rng.choice(
                            sorted(meta_seq[i]-tried_set[i]))
                        tried_set[i].add(candidate[i])
                        # Case (
                        if i in meta_struct.paired_dict:
                            j = meta_struct.paired_dict[i]
                            # DNA Parts
                            if self.part_type == 'DNA':
                                pnt = self.base_compl[candidate[i]][0]
                                candidate[j] = pnt
                            # RNA Parts
                            else:
                                pnt = meta_seq[j]
                                pnt = pnt.intersection(
                                    self.base_compl[candidate[i]])
                                candidate[j] = self.rng.choice(sorted(pnt))
                            tried_set[j] = set(candidate[j])
                    # Case )
                    else:
                        j = meta_struct.rev_paired_dict[i]
                        # DNA Parts
                        if self.part_type == 'DNA':
                            pnt = self.base_compl[candidate[j]][0]
                            candidate[i] = pnt
                        # RNA Parts
                        else:
                            pnt = pnt - tried_set[i]
                            candidate[i] = self.rng.choice(sorted(pnt))
                        tried_set[i].add(candidate[i])

                # Backward phase - Nucleotide choices exhausted, so traceback
                else:
                    # Update roll_back_count
                    roll_back_count += 1

                    # Reset and roll back
                    tried_set[i] = set()
                    kmer_set[i]  = ' '
                    if i > 0:
                        # Clear stuff at current index
                        self._reset_candidate_kmer_set(
                            candidate,
                            kmer_set,
                            i)
                        # Traceback to previous index
                        # since current index done
                        i = self._roll_back(
                            meta_seq,
                            meta_struct,
                            i-1,
                            homology,
                            candidate,
                            tried_set,
                            kmer_set) + 1
                    i -= 1
                    continue

            # See if built candidate[i] is valid
            # Case ( and .
            elif not i in meta_struct.rev_paired_dict:
                # Wrong nucleotide selected
                if not candidate[i] in meta_seq[i]:
                    self._reset_candidate_kmer_set(
                        candidate=candidate,
                        kmer_set=kmer_set,
                        i=i)
                    tried_set[i] = set()
                    continue
            # Case )
            elif i in meta_struct.rev_paired_dict:
                j = meta_struct.rev_paired_dict[i]
                # Paired bases are not complementary
                # or, Wrong nucleotide selected
                if (not candidate[j] in self.base_compl[candidate[i]]) or \
                   (not candidate[i] in meta_seq[i]):
                    self._reset_candidate_kmer_set(
                        candidate=candidate,
                        kmer_set=kmer_set,
                        i=i)
                    tried_set[i] = set()
                    continue

            # See if the local model function is violated
            if local_model_fn:
                rbi = self._get_local_roll_back_index(
                    candidate, i, local_model_fn, verbose)
                # Model function violated, and
                # a traceback location was determined
                if not rbi is None:
                    roll_back_count = 0 if rbi < i else roll_back_count
                    i = self._roll_back(
                        meta_seq,
                        meta_struct,
                        i,
                        homology,
                        candidate,
                        tried_set,
                        kmer_set,
                        rbi)
                    continue

            # Are either of these mers seen previously?
            mmer_seen = False
            kmer = None

            # Handle equal internal and shared repeats
            if i >= homology-1:
                # Get the kmer/rmer
                kmer = ''.join(candidate[i-homology+1:i+1])
                rmer = utils.get_revcomp(kmer)
                mmer = min(kmer, rmer)

                # Case: kmer/rkmer is an internal
                #       repeat to current part
                if not allow_internal_repeat:
                    # Direct repeat
                    if kmer in kmer_set:
                        mmer_seen = True
                    # Inverted repeat
                    elif rmer in kmer_set:
                        mmer_seen = True
                    # Palindrome repeat
                    elif kmer == rmer:
                        mmer_seen = True

                # Case: mmer is a shared repeat with
                #       a previous part
                if not mmer_seen:
                    if mmer in self.kmer_db:
                        mmer_seen = True

                # Case: mmer is a shared repeat with
                #       background
                if not mmer_seen:
                    if not self.background is None:
                        if self.background.K == homology:
                            if mmer in self.background:
                                mmer_seen = True

            # Handle background repeats
            if not mmer_seen and \
               not self.background is None and \
               homology != self.background.K:

                # Determine background K
                K = self.background.K

                # Check is warranted
                if i >= K-1:

                    # Get the kmer/rmer
                    kmer = ''.join(candidate[i-K+1:i+1])
                    rmer = utils.get_revcomp(kmer)
                    mmer = min(kmer, rmer)

                    # Actual check
                    if mmer in self.background:
                        mmer_seen = True

            # Traceback to eliminate repeat
            if mmer_seen:
                i = self._roll_back(
                    meta_seq,
                    meta_struct,
                    i,
                    homology,
                    candidate,
                    tried_set,
                    kmer_set)
                continue

            # Everything OK .. insert kmer
            if i >= homology-1:
                kmer_set[i] = kmer

            # Roll forward
            i += 1

        # Prepare to return candidate
        del kmer_set
        if candidate is None or '-' in candidate:
            return None
        else:
            return ''.join(candidate)

    def _get_opt_pass_count(self, meta_struct, synth_opt, global_model_fn):
        opt_criteria_count = 1 # Since candidate must at least be non-repetitive
        if meta_struct.folding['status']:
            opt_criteria_count += 1
        if synth_opt:
            opt_criteria_count += 1
        if global_model_fn:
            opt_criteria_count += 1
        return opt_criteria_count

    # Diagnostic function -- Shouldn't trigger on experimental changes
    def _is_non_coding_construction_verified(self, meta_seq, meta_struct, candidate):
        i = 0
        while i < len(candidate):
            if not candidate[i] in meta_seq[i]:
                return False,1,i
            if i in meta_struct.paired_dict:
                j = meta_struct.paired_dict[i]
                if not candidate[j] in self.base_compl[candidate[i]]:
                    return False,2,i
            i += 1
        return True,0,0

    def _is_synthesis_verified(self, candidate):
        return self.synthesis.evaluate(candidate)

    def _is_structure_verified(self, meta_struct, struct_type, candidate):
        struct_satisfied = 0

        # Decide which structure criteria to fulfill
        if struct_type == 'centroid':
            candidate_structs = [self.fold.evaluate_centroid(
                candidate)]
        elif struct_type == 'both':
            candidate_structs = [self.fold.evaluate_mfe(
                candidate)]
            candidate_structs.append(self.fold.evaluate_centroid(
                candidate))
        else:
            candidate_structs = [self.fold.evaluate_mfe(
                candidate)]

        for candidate_struct in candidate_structs:
            vienna_meta_struct = self._get_meta_struct(
                candidate_struct)
            # Ensure all forbidden base indices unpaired
            for bp_closed in meta_struct.unpaired_set:
                if not vienna_meta_struct.struct[bp_closed] == '.':
                    return False
            # Ensure all paired base indices paired as desired
            for bp_open in meta_struct.paired_dict:
                if not bp_open in vienna_meta_struct.paired_dict:
                    break
                else:
                    if meta_struct.paired_dict[bp_open] != vienna_meta_struct.paired_dict[bp_open]:
                        break
            else:
                struct_satisfied += 1

        # All structure criteria satisfied
        if struct_satisfied == len(candidate_structs):
            return True
        else:
            return False

    def _get_inverse_fold_candidate(self, start_seq, meta_seq, meta_struct):
        inverse_fold_seq = ''.join(char.lower() if len(meta_seq[i]) == 1 else char for i,char in enumerate(start_seq))
        return self.fold.design(
            seq=inverse_fold_seq, struct=meta_struct.struct).upper()

    def _get_verified_non_coding_candidate(self,
        homology,
        meta_seq,
        meta_struct,
        struct_type,
        synth_opt,
        local_model_fn,
        global_model_fn,
        jump_count,
        fail_count,
        verbose,
        abortion,
        allow_internal_repeat=False):

        # Setup counts and variables
        current_jump_count = 0
        current_fail_count = 0
        struct_fail_count  = 0
        synth_fail_count   = 0
        model_fail_count   = 0
        seed_seq           = None
        verified_candidate = None
        opt_pass_count     = self._get_opt_pass_count(
            meta_struct,
            synth_opt,
            global_model_fn)

        while not verified_candidate:

            # Try to get a non-repetitive candidate
            start_seq = self._get_non_coding_candidate(
                meta_seq=meta_seq,
                meta_struct=meta_struct,
                homology=homology,
                local_model_fn=local_model_fn,
                verbose=verbose,
                jump=current_jump_count < jump_count,
                start_seq=seed_seq,
                allow_internal_repeat=allow_internal_repeat)
            candidate = start_seq
            if start_seq:
                # If structure unmatched but constraint has base pairings, go for inverse-repair strategy
                if meta_struct.inversefolding['status'] and not self._is_structure_verified(
                    meta_struct, struct_type, start_seq):
                    inverse_fold_seq = self._get_inverse_fold_candidate(
                        start_seq,
                        meta_seq,
                        meta_struct)
                    candidate = self._get_non_coding_candidate(
                        meta_seq=meta_seq,
                        meta_struct=meta_struct,
                        homology=homology,
                        local_model_fn=local_model_fn,
                        verbose=verbose,
                        jump=current_jump_count < jump_count,
                        start_seq=inverse_fold_seq,
                        allow_internal_repeat=allow_internal_repeat)

            opt_count = 0

            # If valid candidate then process
            if candidate:
                current_jump_count = 0
                opt_count += 1

                # Diagnostic block -- Shouldn't trigger on experimental changes
                construction = self._is_non_coding_construction_verified(
                    meta_seq,
                    meta_struct,
                    candidate)
                construction_state = construction[0]
                error_digest = construction[1:]
                if not construction_state:
                    current_fail_count += 1
                    raise Exception(
                        'Maker built a rogue candidate: {}\nError Digest: {}\nPlease report issue to authors.'.format(
                            candidate,
                            error_digest))
                else:
                    pass

                # Synthesis optimization
                if synth_opt:
                    if self._is_synthesis_verified(candidate):
                        opt_count += 1
                    else:
                        synth_fail_count += 1

                # Global model optimization
                if global_model_fn:

                    # Try to evaluate the global_model_fn on candidate_str
                    outcome = True
                    try:
                        outcome = global_model_fn(candidate)
                    except Exception as e:
                        print(' Global Model fn. failed to evaluate complete path: {}\n'.format(
                            candidate))
                        raise e

                    # Outcome satisfactory?
                    try:
                        assert outcome in [True, False]
                    except Exception as e:
                        print(' Global Model fn. returned a non-boolean state: {}\n'.format(
                            outcome))
                        raise e

                    # Process outcome
                    if outcome: # True
                        opt_count += 1
                    else:       # False
                        model_fail_count += 1

                # Structural optimization
                if meta_struct.folding['status']:
                    if self._is_structure_verified(
                        meta_struct,
                        struct_type,
                        candidate):
                        opt_count += 1
                    else:
                        struct_fail_count += 1

                # Did everything get optimized?
                if opt_count == opt_pass_count:
                    verified_candidate = candidate
                else:
                    current_fail_count += 1

                # Failure count exceeded, terminate
                if current_fail_count == fail_count:
                    break

            # No candidate produced
            else:
                # No jumps made yet no non-repetitive candidate found
                if current_jump_count >= jump_count:
                    break
                # Increase current_jump_count
                else:
                    current_jump_count += 1
                # Abortion limit reached?
                if abortion and current_jump_count >= jump_count:
                    break

        if int(verbose) > 1:
            print('\n  [seq fails] {}, [struct fails] {}, [synth fails] {}, [global fails] {}, [opt fails] {}'.format(
                current_jump_count,
                struct_fail_count,
                synth_fail_count,
                model_fail_count,
                current_fail_count))

        return verified_candidate, current_jump_count+1, current_fail_count+1

    def _get_non_coding_nrps(
        self,
        homology,
        seq,
        struct,
        struct_type,
        target,
        synth_opt,
        local_model_fn,
        global_model_fn,
        jump_count,
        fail_count,
        verbose,
        abortion,
        allow_internal_repeat=False):

        # Setup structures
        seq = seq.upper()
        meta_struct = self._get_meta_struct(
            struct=self._get_adjusted_struct(struct, seq))
        meta_seq = self._get_meta_seq(
            seq=seq,
            meta_struct=meta_struct)

        seq_count  = 0
        iter_count = 0
        time_sum   = 0.0
        begin_time = time()
        break_flag = False

        # Setup Bernoulli Success model
        total_jump_trials    = jump_count
        total_jump_successes = 1
        curr_jump_prob       = berno.get_prob(
            trials=total_jump_trials,
            success=total_jump_successes)
        curr_jump_trial      = jump_count
        total_fail_trials    = fail_count
        total_fail_successes = 1
        curr_fail_prob       = berno.get_prob(
            trials=total_fail_trials,
            success=total_fail_successes)
        curr_fail_trial      = fail_count

        # Stream parts until completion
        while True:

            t0 = time()

            candidate, curr_jump_trial, curr_fail_trial = self._get_verified_non_coding_candidate(
                homology,
                meta_seq,
                meta_struct,
                struct_type,
                synth_opt,
                local_model_fn,
                global_model_fn,
                curr_jump_trial,
                curr_fail_trial,
                verbose,
                abortion,
                allow_internal_repeat)

            if candidate is None:
                break_flag = True

            # Got a canidate -- will try again
            if not break_flag:
                final_candidate = candidate
                update_status   = True
                try:
                    for kmer in utils.stream_min_kmers(
                        seq=candidate,
                        k=homology):
                        self.kmer_db.add(kmer)
                except Exception as E:
                    update_status = False

                if not update_status: # Memory full
                    if verbose:
                        print('[ERROR] Memory Full ... Breaking Loop')
                        yield update_status

                seq_count  += 1
                time_sum   += time()-t0
                iter_count += 1

                if verbose:
                    print(' [part] {}, [{}-mers] {}, [iter time] {:.2f}s, [avg time] {:.2f}s, [total time] {:.2f}h'.format(
                        seq_count,
                        homology,
                        len(self.kmer_db),
                        time()-t0, time_sum / iter_count,
                        (time() - begin_time) / 3600.0))

                yield final_candidate

                # No more parts required
                if seq_count == target:
                    break
            # No more candidates to build
            else:
                yield candidate
                break

            # Update failure limits based on Bernoulli Success model
            total_jump_trials    += curr_jump_trial
            total_jump_successes += 1
            curr_jump_prob       =  berno.get_prob(trials=total_jump_trials, success=total_jump_successes)
            curr_jump_trial      =  berno.get_trials(prob=curr_jump_prob)
            total_fail_trials    += curr_fail_trial
            total_fail_successes += 1
            curr_fail_prob       =  berno.get_prob(trials=total_fail_trials, success=total_fail_successes)
            curr_fail_trial      =  berno.get_trials(prob=curr_fail_prob)

    def _check_maker_constraints(
        self,
        seq,
        struct,
        part_type,
        allow_internal_repeat,
        target,
        homology):
        # Sequence Legality 1
        if not isinstance(seq, str):
            print('\n [ERROR]    Sequence Constraint must be a string, not {}'.format(type(seq)))
            print(' [SOLUTION] Try correcting Sequence Constraint\n')
            return False
        # Sequence Legality 2
        if len(seq) < 5:
            print('\n [ERROR]    Sequence Constraint must be longer than 4 bases, not {}'.format(len(seq)))
            print(' [SOLUTION] Try using a longer Sequence Constraint\n')
            return False
        # Structure Legality 1
        if not isinstance(struct, str):
            print('\n [ERROR]    Structure Constraint must be a string, not {}'.format(type(struct)))
            print(' [SOLUTION] Try correcting Structure Constraint\n')
            return False
        # Structure Legality 2
        if len(struct) != len(seq):
            print('\n [ERROR]    Structure Constraint must be same length as Sequence Constraint ({}), not {}'.format(
                len(seq),
                len(struct)))
            print(' [SOLUTION] Try correcting length of Structure Constraint\n')
            return False
        # Part Type Legality 1
        if not isinstance(part_type, str):
            print('\n [ERROR]    Part Type must be a string, not \'{}\''.format(type(part_type)))
            print(' [SOLUTION] Try correcting Part Type\n')
            return False
        # Part Type Legality 2
        if not part_type in ['DNA', 'RNA']:
            print('\n [ERROR]    Part Type must be \'RNA\' or \'DNA\', not \'{}\''.format(part_type))
            print(' [SOLUTION] Try correcting Part Type\n')
            return False
        # Sequence Legality 3
        seq_legal, seq_illegal_chars = makerchecks.is_seq_constr_legal(seq, part_type)
        if not seq_legal:
            print('\n [ERROR]    {} Sequence Constraint is not legal due to chars: {}'.format(part_type, seq_illegal_chars))
            print(' [SOLUTION] Try correcting Sequence Constraint or Part Type\n')
            return False
        # Structure Legality 3
        struct_legal, unclosed, unopened, invalid = makerchecks.is_structure_valid(struct)
        if not struct_legal:
            print('\n [ERROR]    Structure Constraint is illegal or unbalanced')
            if unclosed:
                print(' [ERROR]    >> Unclosed bases at locations: {}'.format(unclosed))
            if unopened:
                print(' [ERROR]    >> Unopened bases at locations: {}'.format(unopened))
            if invalid:
                print(' [ERROR]    >> Invalid characters at locations: {}'.format(invalid))
            print(' [SOLUTION] Try correcting Structure Constraint\n')
            return False
        # Sequence + Structure + Part Type Base Pairing Combination Legality
        combo_state, incompat_locs, reduced_locs = makerchecks.is_pairing_compatible(seq, struct, part_type)
        if not combo_state:
            if incompat_locs:
                print('\n [ERROR]    Incompatible Base Pairing for {} Parts at locations: {}'.format(part_type, incompat_locs))
                print(' [SOLUTION] Try correcting Sequence Constraint or Part Type\n')
                return False
            if reduced_locs: # -- soft check
                print('\n [WARNING]  Reducible Paired Bases at locations: {}'.format(reduced_locs))
                print(' [WARNING]  Fewer Parts may be Generated')
        # Lmax Legality 1
        if not isinstance(homology, int):
            print('\n [ERROR]    Lmax must be an integer, not {}'.format(type(homology)))
            print(' [SOLUTION] Try correcting Lmax\n')
            return False
        # Lmax Legality 2
        if homology-1 < 5:
            print('\n [ERROR]    Lmax must be greater than 4, not {}'.format(homology-1))
            print(' [SOLUTION] Try correcting Lmax\n')
            return False
        # Lmax Legality 3
        if homology-1 >= len(seq):
            print('\n [ERROR]    Lmax must be less than length of Sequence Constraint ({}), not {}'.format(len(seq), homology-1))
            print(' [SOLUTION] Try correcting Lmax\n')
            return False
        # Target Size Legality 1
        if not isinstance(target, int):
            print('\n [ERROR]    Target Size must be an integer, not {}'.format(type(target)))
            print(' [SOLUTION] Try correcting Target Size\n')
            return False
        # Target Size Legality 2
        if target < 1:
            print('\n [ERROR]    Target Size must be greater than 0, not {}'.format(target))
            print(' [SOLUTION] Try correcting Target Size\n')
            return False
        # Sequence Sufficiency -- soft check
        seq_sufficient, constrained_motif_locs = makerchecks.is_seq_constr_sufficient(seq, struct, homology, target)
        if not seq_sufficient:
            print('\n [WARNING]  Target Size of {} may be unreachable from given Sequence/Structure Constraint and Lmax of {}'.format(target, homology-1))
            print(' [WARNING]  >> Lmax limiting windows between locations: {}'.format(constrained_motif_locs))
            print(' [WARNING]  Fewer Parts may be Generated')
        # Internal Repeats Legality
        if not allow_internal_repeat in [True, False]:
            print('\n [ERROR]    Internal Repeat must be boolean, not {}'.format(type(allow_internal_repeat)))
            print(' [SOLUTION] Try correcting Internal Repeat\n')
            return False
        # Structure Sufficiency
        if not allow_internal_repeat:
            struct_sufficient, long_hairpins = makerchecks.is_structure_not_conflict(struct, homology)
            if not struct_sufficient:
                print('\n [ERROR] Structure Constraint is insufficient based on given Lmax')
                for long_hairpin in long_hairpins:
                    print(' [ERROR] >> Long hairpin at locations: {}'.format(long_hairpin))
                    print(' [SOLUTION] Try relaxing Structure Constraint or setting internal_repeats=True')
                print()
                return False
        return True

    def _check_maker_inputs(
        self,
        struct_type,
        synth_opt,
        jump_count,
        fail_count,
        output_file,
        verbose):
        # Struct Type Legality
        if not struct_type in ['mfe', 'centroid', 'both']:
            print('\n [ERROR]    Struct Type must be \'mfe\', \'centroid\' or \'both\', not \'{}\''.format(struct_type))
            print(' [SOLUTION] Try correcting Part Type\n')
            return False
        # Synth Optimization Legality
        if not synth_opt in [True, False]:
            print('\n [ERROR]    Synth Opt must be True or False, not {}'.format(synth_opt))
            print(' [SOLUTION] Try correcting Synth Opt\n')
            return False
        # Jump Count Legality 1
        if not isinstance(jump_count, int):
            print('\n [ERROR]    Jump Count must be an integer, not {}'.format(type(jump_count)))
            print(' [SOLUTION] Try correcting Jump Count\n')
            return False
        # Jump Count Legality 2
        if jump_count < 0:
            print('\n [ERROR]    Jump Count must be greater than 0, not {}'.format(jump_count))
            print(' [SOLUTION] Try correcting Jump Count\n')
            return False
        # Fail Count Legality 1
        if not isinstance(fail_count, int):
            print('\n [ERROR]    Fail Count must be an integer, not {}'.format(type(fail_count)))
            print(' [SOLUTION] Try correcting Fail Count\n')
            return False
        # Fail Count Legality 2
        if fail_count < 0:
            print('\n [ERROR]    Fail Count must be greater than 0, not {}'.format(fail_count))
            print(' [SOLUTION] Try correcting Fail Count\n')
            return False
        # Output File Legality
        if not output_file is None:
            if not isinstance(output_file, str):
                print('\n [ERROR]    Output File must be a string or None, not {}'.format(type(output_file)))
                print(' [SOLUTION] Try correcting Output File\n')
                return False
        # Everything OK
        return True

    def nrp_maker(self,
        homology,
        seq_constr,
        struct_constr,
        target_size,
        background=None,
        struct_type=None,
        synth_opt=True,
        local_model_fn=None,
        global_model_fn=None,
        jump_count=10,
        fail_count=1000,
        output_file=None,
        verbose=True,
        abortion=True,
        allow_internal_repeat=False,
        check_constraints=True):

        if verbose:
            print('\n[Non-Repetitive Parts Calculator - Maker Mode]')
        build_parts = True

        # Check Maker Constraints
        if check_constraints:
            if verbose:
                print('\n[Checking Constraints]')
                print('  Sequence Constraint: {}'.format(seq_constr))
                print(' Structure Constraint: {}'.format(struct_constr))
                print('      Part Type      : {}'.format(self.part_type))
                print('           Lmax      : {} bp'.format(homology-1))
                print('    Target Size      : {} parts'.format(target_size))
                print('  Internal Repeats   : {}'.format(allow_internal_repeat))
            check_status = self._check_maker_constraints(
                seq_constr,
                struct_constr,
                self.part_type,
                allow_internal_repeat,
                target_size,
                homology)

            if check_status == False:
                if verbose:
                    print(' Check Status: FAIL\n')
                build_parts = False
            else:
                if verbose:
                    print('\n Check Status: PASS')

            # Background Check
            if build_parts:
                if not background is None:
                    if verbose:
                        print('\n[Checking Background]\n Background: {}'.format(background))
                    if isinstance(background, kmerSetDB):
                        if background.K > len(seq_constr):
                            build_parts = False
                            print('\n [ERROR]    Background Lmax of {} is greater than desired part length ({}-bp)'.format(
                                background.K-1,
                                len(seq_constr)))
                            print(' [SOLUTION] Try using a background with Lmax less than or equal to part length\n')
                            if verbose:
                                print(' Check Status: FAIL\n')
                        if build_parts and not background.ALIVE:
                            build_parts = False
                            print('\n [ERROR]    Background is closed or dropped')
                            print(' [SOLUTION] Try using an open Background\n')
                            if verbose:
                                print(' Check Status: FAIL\n')
                        if build_parts:
                            if verbose:
                                print('\n Check Status: PASS')

                    else:
                        build_parts = False
                        print('\n [ERROR]    Background Object is INVALID')
                        print(' [SOLUTION] Try instantiating background via nrpcalc.background(...)\n')
                        if verbose:
                            print(' Check Status : FAIL\n')

            # Arguments Check
            if build_parts:
                if verbose:
                    print('\n[Checking Arguments]')
                    print(' Struct Type : {}'.format(struct_type))
                    print('  Synth Opt  : {}'.format(synth_opt))
                    print('   Jump Count: {}'.format(jump_count))
                    print('   Fail Count: {}'.format(fail_count))
                    print(' Output File : {}'.format(output_file))
                check_status = self._check_maker_inputs(
                    struct_type,
                    synth_opt,
                    jump_count,
                    fail_count,
                    output_file,
                    verbose)

                if check_status == False:
                    if verbose:
                        print(' Check Status: FAIL\n')
                    build_parts = False
                else:
                    if verbose:
                        print('\n Check Status: PASS')

            if not build_parts:
                # Cleanups
                self.background = None
                self.kmer_db = None
                raise RuntimeError('Invalid Constraints, Background or Arguments')

        # Separate Checks from Build Logs
        if verbose:
            print()

        # kmer_db and background Setup
        projector.setup_proj_dir(self.proj_id)
        self.kmer_db = set()
        self.background = background
        # self.kmer_db = kmerSetDB(
        #     path='./{}/kmerDB'.format(self.proj_id),
        #     homology=homology,
        #     verbose=verbose)

        # Project Setup
        if output_file is None:
            projector.setup_proj_dir(self.proj_id)
            output_file = './{}/seq_list.fa'.format(self.proj_id)

        # Execute Maker
        with open(output_file, 'w') as out_file:
            seq_constr    = seq_constr.upper()
            struct_constr = self._get_adjusted_struct(
                struct_constr,
                seq_constr)
            current_nrp_count = 0
            memory_exhausted  = False

            if verbose:
                print('Constructing Toolbox:\n')

            for non_coding_nrp in self._get_non_coding_nrps(
                homology,
                seq_constr,
                struct_constr,
                struct_type,
                target_size,
                synth_opt,
                local_model_fn,
                global_model_fn,
                jump_count,
                fail_count,
                verbose,
                abortion,
                allow_internal_repeat):

                # Write out genetic part
                if non_coding_nrp:
                    out_file.write(
                        '>non-repetitive part {}\n'.format(
                            current_nrp_count+1))
                    non_coding_nrp = '\n'.join(
                        textwrap.wrap(
                            non_coding_nrp, 80))
                    out_file.write('{}\n'.format(
                        non_coding_nrp))
                    current_nrp_count += 1
                else:
                    if non_coding_nrp is None:
                        if verbose:
                            print('Failure Limits Exceeded or k-mers Exhausted. Cannot Build More Parts.')
                    else:
                        if verbose:
                            print('Memory Capacity at Full. Cannot Build More Parts.')
                        memory_exhausted = True

                # Memory no longer available ... stop
                if memory_exhausted:
                    break
            if verbose:
                print('\nConstruction Complete.\n')

        # Detach Background
        self.background = None

        # Remove kmerSetDB
        # self.kmer_db.drop()
        self.kmer_db = None

        # Pack output in dictionary
        parts_dict = {}
        for i,line in enumerate(utils.stream_fasta_seq_list(output_file)):
            line = line.strip()
            parts_dict[i] = line

        # Cleanups and Return
        projector.remove_proj_dir()
        if verbose:
            print('Non-Repetitive Toolbox Size: {}'.format(current_nrp_count))
        return parts_dict

def main():
    homology = 16
    sm_obj = NRPMaker(seed=7)

    # Hammerhead Nielsen Paper
    seq         = 'NNNN AGNNNU CANNNNN UGUGCUU NNNNNU CUGAUGA NNNN GUGA NNNN GAAA NNNC CUCU NNNNN UAAU NNNNN UUAA NNNN' # Nielsen Like
    struct      = 'xxxx x((((( x(((((( xxxxxxx )))))) xxxxxxx (((( xxxx )))) xxx) )))) xxxx ((((( xxxx ))))) xxxx xxxx'
    seq         = ''.join(seq.split(' '))
    struct      = ''.join(struct.split(' '))
    output_file = None #'riboz.fa'
    background  = None #utils.get_fasta_seq_list(fasta_filename='input.fa.bk104')

    # Final Result Store
    final_toolbox = {}

    # Initialize Background
    background = kmerSetDB(
        path='./testDB',
        homology=homology,
        verbose=True)
    # background.multiadd(
    #     utils.get_fasta_seq_list(
    #         fasta_filename='input.fa.bk104'))

    # Background Based Single Part Design Works
    t0 = time()
    tt = 0
    toolbox1 = sm_obj.nrp_maker(homology, [seq], [struct], [10],
        struct_type='mfe',
        background=background,
        jump_count=100,
        fail_count=1000,
        synth_opt=False,
        verbose=True,
        abortion=True,
        allow_internal_repeat=True,
        output_file=output_file)
    tt += time() - t0
    final_toolbox.update(
        zip(range(len(final_toolbox), len(final_toolbox)+len(toolbox1)),
            toolbox1.values()))

    # Adding More Background Post Design Works
    background.multiadd(toolbox1.values())

    # Serial Designs from Multiple Constraints Works
    t0 = time()
    toolbox2 = sm_obj.nrp_maker(homology, [seq]*2, [struct]*2, [20]*2,
        struct_type='mfe',
        background=background,
        jump_count=100,
        fail_count=1000,
        synth_opt=False,
        verbose=True,
        abortion=True,
        allow_internal_repeat=False,
        output_file=output_file)
    tt += time() - t0
    final_toolbox.update(
        zip(range(len(final_toolbox), len(final_toolbox)+len(toolbox2)),
            toolbox2.values()))

    # Assert All Parts Unique and Non-Repetitive
    assert len(set(final_toolbox.values())) == len(
        finder.nrp_finder(final_toolbox.values(), homology, None, verbose=False))

    # Drop Background
    background.drop()

    # Report Time Elapsed
    print('\nWall Time {} sec'.format(tt))

if __name__ == '__main__':
    main()


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

import numpy as np

from .utils import geomutils
from .utils import obutils

class HydrateMoleculeLegacy:
    def __init__(self, distance=3.0, charge=0, atom_type="W"):
        """Initialize the legacy hydrate typer for AutoDock 4.2.x

        Args:
            distance (float): distance between water molecules and ligand heavy atoms. (default: 3.0)
            charge (float): partial charge of the water molecule. Not use for the hydrated docking. (default: 0)
            atom_type (str): atom type of the water molecule. (default: W)

        """
        self._distance = distance
        self._charge = charge
        self._atom_type = atom_type
        self._bond_type = 1
        self._rotatable = False
        self._hb_config = {'HD': {1: (1, 1)},    # neigh: 1, wat: 1, sp1
                           'OA': {
                                      1: (2, 2), # neigh: 1, wat: 2, sp2
                                      2: (2, 3)  # neigh: 2, wat: 2, sp3
                                  },
                           'SA': {
                                      1: (2, 2), # neigh: 1, wat: 2, sp2
                                      2: (2, 3)  # neigh: 2, wat: 2, sp3
                                  },
                           'NA': {
                                      1: (1, 1), # neigh: 1, wat: 3, sp1
                                      2: (1, 2), # neigh: 2, wat: 1, sp2
                                      3: (1, 3)  # neigh: 3, wat: 1, sp3
                                  }
                           }

    def _place_sp1_one_water(self, anchor_xyz, neighbor_xyz, hb_length=3.0):
        position = anchor_xyz + geomutils.vector(neighbor_xyz, anchor_xyz)
        position = geomutils.resize_vector(position, hb_length, anchor_xyz)
        positions = np.array([position])

        return positions

    def _place_sp2_one_water(self, anchor_xyz, neighbor1_xyz, neighbor2_xyz, hb_length=3.0):
        position = geomutils.atom_to_move(anchor_xyz, [neighbor1_xyz, neighbor2_xyz])
        position = geomutils.resize_vector(position, hb_length, anchor_xyz)
        positions = np.array([position])

        return positions

    def _place_sp2_two_waters(self, anchor_xyz, neighbor1_xyz, neighbor2_xyz, hb_lengths, angles):
        if len(hb_lengths) != 2:
            raise ValueError()
        if len(angles) != 2:
            raise ValueError()

        positions = []

        r = geomutils.rotation_axis(neighbor1_xyz, anchor_xyz, neighbor2_xyz, origin=anchor_xyz)
        p = neighbor1_xyz

        # We rotate p to get each vectors if necessary
        for hb_length, angle in zip(hb_lengths, angles):
            vector = p
            if angle != 0.:
                position = geomutils.rotate_point(vector, anchor_xyz, r, angle)
            position = geomutils.resize_vector(position, hb_length, anchor_xyz)
            positions.append(position)

        positions = np.array(positions)

        return positions

    def _place_sp3_one_water(self, anchor_xyz, neighbor1_xyz, neighbor2_xyz, neighbor3_xyz, hb_length):
        # We have to normalize bonds, otherwise the water molecule is not well placed
        v1 = anchor_xyz + geomutils.normalize(geomutils.vector(anchor_xyz, neighbor1_xyz))
        v2 = anchor_xyz + geomutils.normalize(geomutils.vector(anchor_xyz, neighbor2_xyz))
        v3 = anchor_xyz + geomutils.normalize(geomutils.vector(anchor_xyz, neighbor3_xyz))

        position = geomutils.atom_to_move(anchor_xyz, [v1, v2, v3])
        position = geomutils.resize_vector(position, hb_length, anchor_xyz)
        positions = np.array([position])

        return positions

    def _place_sp3_two_waters(self, anchor_xyz, neighbor1_xyz, neighbor2_xyz, hb_lengths, angles):
        if len(hb_lengths) != 2:
            raise ValueError()
        if len(angles) != 2:
            raise ValueError()

        positions = []

        v1 = anchor_xyz + geomutils.normalize(geomutils.vector(anchor_xyz, neighbor1_xyz))
        v2 = anchor_xyz + geomutils.normalize(geomutils.vector(anchor_xyz, neighbor2_xyz))

        r = anchor_xyz + geomutils.normalize(geomutils.vector(v1, v2))
        p = geomutils.atom_to_move(anchor_xyz, [v1, v2])

        # We rotate p to get each vectors if necessary
        for hb_length, angle in zip(hb_lengths, angles):
            vector = p
            if angle != 0.:
                position = geomutils.rotate_point(vector, anchor_xyz, r, angle)
            position = geomutils.resize_vector(position, hb_length, anchor_xyz)
            positions.append(position)

        positions = np.array(positions)

        return positions

    def hydrate(self, mol):
        """Add water molecules to the ligand

        Args:
            mol (OBMol): input OBMol molecule object

        """
        setup = mol.setup
        water_anchors = []
        water_positions = []
        # It will be the same distance for all of the water molecules
        hb_length = self._distance

        for a, neighbors in setup.graph.items():
            atom_type = setup.get_atom_type(a)
            anchor_xyz = setup.get_coord(a)
            neighbor1_xyz = setup.get_coord(neighbors[0])
            positions = np.array([])
            n_wat = None
            hyb = None

            if atom_type in self._hb_config:
                try:
                    n_wat, hyb = self._hb_config[atom_type][len(neighbors)]
                except KeyError:
                    raise RuntimeError('Cannot place water molecules on atom %d of type %s with %d neighbors.' % (a, atom_type, len(neighbors)))

                water_anchors.append(a)

            if hyb == 1:
                if n_wat == 1:
                    # Example: X-HD
                    positions = self._place_sp1_one_water(anchor_xyz,
                                                          neighbor1_xyz,
                                                          hb_length - 1.0)
            elif hyb == 2:
                if n_wat == 1:
                    # Example: X-Nitrogen-X
                    neighbor2_xyz = setup.get_coord(neighbors[1])
                    positions = self._place_sp2_one_water(anchor_xyz,
                                                          neighbor1_xyz, neighbor2_xyz,
                                                          hb_length)
                elif n_wat == 2:
                    # Example: C=0 (backbone oxygen)
                    tmp_neighbors = [x for x in setup.get_neigh(neighbors[0]) if not x == a]
                    neighbor2_xyz =  setup.get_coord(tmp_neighbors[0])
                    positions = self._place_sp2_two_waters(anchor_xyz,
                                                           neighbor1_xyz, neighbor2_xyz,
                                                           [hb_length, hb_length],
                                                           [-np.radians(120), np.radians(120)])
                elif n_wat == 3:
                    hyb = 3
            elif hyb == 3:
                if n_wat == 1:
                    # Example: Ammonia
                    neighbor2_xyz = setup.get_coord(neighbors[1])
                    neighbor3_xyz = setup.get_coord(neighbors[2])
                    positions = self._place_sp3_one_water(anchor_xyz,
                                                          neighbor1_xyz, neighbor2_xyz, neighbor3_xyz,
                                                          hb_length)
                elif n_wat == 2:
                    # Example: O-HD (Oxygen in hydroxyl group)
                    neighbor2_xyz = setup.get_coord(neighbors[1])
                    positions = self._place_sp3_two_waters(anchor_xyz,
                                                           neighbor1_xyz, neighbor2_xyz,
                                                           [hb_length, hb_length],
                                                           [-np.radians(60), np.radians(60)])
                elif n_wat == 3:
                    positions = np.array([])

            if positions.size:
                water_positions.append(positions)

        for water_anchor, waters_on_anchor in zip(water_anchors, water_positions):
            for water_on_anchor in waters_on_anchor:
                tmp = setup.pdbinfo[water_anchor]
                pdbinfo = obutils.PDBAtomInfo('WAT', tmp.resName, tmp.resNum, tmp.chain)
                setup.add_pseudo(water_on_anchor, self._charge, [water_anchor], self._atom_type,
                                     self._bond_type, self._rotatable, pdbinfo)


#############MEASURING FACIAL SIMILARITY USING SIAMESE NETWORK##########
#############Created & Modified by: MUHAMMAD UMER#######################
#############DATE: 11/5/2020############################################

import torchvision
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.utils.data import DataLoader,Dataset
import matplotlib.pyplot as plt
import torchvision.utils
import numpy as np
import random
from PIL import Image
import torch
from torch.autograd import Variable
import PIL.ImageOps
import torch.nn as nn
from torch import optim
import torch.nn.functional as F
from Dataset import Config, SiameseNetworkDataset
from Model import  SiameseNetwork, ContrastiveLoss

def imshow(img,text=None,should_save=False):
    npimg = img.numpy()
    plt.axis("off")
    if text:
        plt.text(75, 8, text, style='italic',fontweight='bold',
            bbox={'facecolor':'white', 'alpha':0.8, 'pad':10})
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

def show_plot(iteration,loss):
    plt.plot(iteration,loss)
    plt.show()

################JUST PLOTTING RANDOM FACES TO SEE###############################
folder_dataset = dset.ImageFolder(root=Config.training_dir)
siamese_dataset = SiameseNetworkDataset(imageFolderDataset=folder_dataset,
                                        transform=transforms.Compose([transforms.Resize((100,100)),
                                                                      transforms.ToTensor()
                                                                      ])
                                       ,should_invert=False)

vis_dataloader = DataLoader(siamese_dataset,
                        shuffle=True,
                        num_workers=8,
                        batch_size=8)
dataiter = iter(vis_dataloader)

example_batch = next(dataiter)
concatenated = torch.cat((example_batch[0],example_batch[1]),0)
imshow(torchvision.utils.make_grid(concatenated))
print(example_batch[2].numpy())

#####################################TRAINING PHASE##########################################

train_dataloader = DataLoader(siamese_dataset,
                        shuffle=True,
                        num_workers=8,
                        batch_size=Config.train_batch_size)

net = SiameseNetwork().cuda()
criterion = ContrastiveLoss()
optimizer = optim.Adam(net.parameters(), lr=0.0005)
counter = []
loss_history = []
iteration_number = 0

for epoch in range(0,Config.train_number_epochs):
    for i, data in enumerate(train_dataloader, 0):
        img0, img1, label = data
        img0, img1, label = img0.cuda(), img1.cuda(), label.cuda()
        optimizer.zero_grad()
        output1, output2 = net(img0, img1)
        loss_contrastive = criterion(output1,output2,label)
        loss_contrastive.backward()
        optimizer.step()
        if i %10 == 0 :
            print("Epoch number {}\n Current loss {}\n".format(epoch,loss_contrastive.item()))
            iteration_number +=10
            counter.append(iteration_number)
            loss_history.append(loss_contrastive.item())
show_plot(counter,loss_history)

#####################################TESTING PHASE##########################################

folder_dataset_test = dset.ImageFolder(root=Config.testing_dir)
siamese_dataset = SiameseNetworkDataset(imageFolderDataset=folder_dataset_test,
                                        transform=transforms.Compose([transforms.Resize((100, 100)),
                                                                      transforms.ToTensor()
                                                                      ])
                                        , should_invert=False)

test_dataloader = DataLoader(siamese_dataset, num_workers=6, batch_size=1, shuffle=True)
dataiter = iter(test_dataloader)
x0, _, _ = next(dataiter)

for i in range(10):
    _, x1, label2 = next(dataiter)
    concatenated = torch.cat((x0, x1), 0)

    output1, output2 = net(Variable(x0).cuda(), Variable(x1).cuda())
    euclidean_distance = F.pairwise_distance(output1, output2)
    imshow(torchvision.utils.make_grid(concatenated), 'Dissimilarity: {:.2f}'.format(euclidean_distance.item()))