"""
policy

Created by: <NAME>
On: 22-11-19, 13:30
"""
import numpy as np
import torch
from abc import ABC, abstractmethod
from torch.utils.data import TensorDataset, DataLoader
from tqdm import trange

from drugex.api.environ.models import Environ
from drugex.api.pretrain.generators import Generator, PolicyAwareGenerator
from drugex.core import util


class PolicyGradient(ABC):

    def __init__(self, batch_size=512, mc=10, epsilon=0.01, beta=0.1):
        self.batch_size = batch_size
        self.mc = mc
        self.epsilon = epsilon
        self.beta = beta

    @abstractmethod
    def __call__(self, environ : Environ, exploit : Generator, explore = None):
        pass


class PG(PolicyGradient):

    def __call__(self, environ: Environ, exploit: PolicyAwareGenerator, explore=None):
        """Training generator under reinforcement learning framework,
        The rewoard is only the final reward given by environment (predictor).

        agent (model.Generator): the exploitation network for SMILES string generation
        environ (util.Activity): the environment provide the final reward for each SMILES
        explore (model.Generator): the exploration network for SMILES string generation,
            it has the same architecture with the agent.
        """

        smiles, valids, seqs = exploit.sample(
            self.batch_size
            , explore=explore
            , epsilon=self.epsilon
            , include_tensors=True
            , mc=self.mc
        )

        # obtaining the reward
        preds = environ.predictSMILES(smiles)
        preds[valids == False] = 0
        preds -= self.beta
        preds = torch.Tensor(preds.reshape(-1, 1)).to(util.getDev())

        ds = TensorDataset(seqs, preds)
        loader = DataLoader(ds, batch_size=self.batch_size)

        # Training Loop
        for seq, pred in loader:
            exploit.policyUpdate(seq, pred)


def rollout_pg(agent, environ, explore=None, *, batch_size, baseline, mc, epsilon):
    """Training generator under reinforcement learning framework.

    The reward is given for each token in the SMILES, which is generated by
    Monte Carlo Tree Search based on final reward given by the environment.

    Arguments:

        agent (model.Generator): the exploitation network for SMILES string generation
        environ (util.Activity): the environment provide the final reward for each SMILES
        explore (model.Generator): the exploration network for SMILES string generation,
            it has the same architecture with the agent.
    """

    agent.optim.zero_grad()
    seqs = agent.sample(batch_size, explore=explore, epsilon=epsilon)
    batch_size = seqs.size(0)
    seq_len = seqs.size(1)
    rewards = np.zeros((batch_size, seq_len))
    smiles, valids = util.check_smiles(seqs, agent.voc)
    preds = environ(smiles) - baseline
    preds[valids == False] = - baseline
    scores, hiddens = agent.likelihood(seqs)

    # Monte Carlo Tree Search for step rewards generation
    for _ in trange(mc):
        for i in range(seq_len):
            if (seqs[:, i] != 0).any():
                h = hiddens[:, :, i, :]
                subseqs = agent.sample(batch_size, inits=(seqs[:, i], h, i + 1, None))
                subseqs = torch.cat([seqs[:, :i+1], subseqs], dim=1)
                subsmile, subvalid = util.check_smiles(subseqs, voc=agent.voc)
                subpred = environ(subsmile) - baseline
                subpred[1 - subvalid] = -baseline
            else:
                subpred = preds
            rewards[:, i] += subpred
    loss = agent.PGLoss(scores, seqs, torch.FloatTensor(rewards / mc))
    loss.backward()
    agent.optim.step()
    return 0, valids.mean(), smiles, preds