from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from .data import Data
from .. import config
from ..backend import tf
from ..utils import get_num_args, run_if_all_none


class PDE(Data):
    """ODE or time-independent PDE solver.

    Args:
        geometry: Instance of ``Geometry``.
        pde: A global PDE or a list of PDEs. ``None`` if no global PDE.
        bcs: A boundary condition or a list of boundary conditions. ``[]`` if no boundary condition.
        num_domain (int): The number of training residual points sampled inside the domain.
        num_boundary (int): The number of training residual points sampled on the boundary.
        train_distribution (string): The distribution to sample training residual points. One of the following:
            "uniform" (equispaced grid), "pseudo" (pseudorandom) or "sobol" (Sobol sequence).
        anchors: A Numpy array of training residual points, in addition to the `num_domain` and `num_boundary` sampled
            points.
        solution: The reference solution.
        num_test: The number of residual points for testing the PDE residual.

    Attributes:
        train_x_all: A Numpy array of all residual points for training. `train_x_all` is unordered,
            and does not have duplication.
        train_x: A Numpy array of the residual points fed into the network for training.
            `train_x` is a subset of `train_x_all`, ordered from BCs to PDE, and may have duplicate points.
        num_bcs (list): `num_bcs[i]` is the number of residual points for `bcs[i]`.
        test_x: A Numpy array of the residual points fed into the network for testing the PDE residual.
    """

    def __init__(
        self,
        geometry,
        pde,
        bcs,
        num_domain=0,
        num_boundary=0,
        train_distribution="sobol",
        anchors=None,
        solution=None,
        num_test=None,
    ):
        self.geom = geometry
        self.pde = pde
        self.bcs = bcs if isinstance(bcs, (list, tuple)) else [bcs]

        self.num_domain = num_domain
        self.num_boundary = num_boundary
        self.train_distribution = train_distribution
        self.anchors = anchors

        self.soln = solution
        self.num_test = num_test

        self.train_x_all = None
        self.train_x, self.train_y = None, None
        self.num_bcs = None
        self.test_x, self.test_y = None, None

        self._check()
        self.train_next_batch()
        self.test()

    def _check(self):
        if self.train_distribution not in ["uniform", "pseudo", "sobol"]:
            raise ValueError(
                "self.train_distribution == {}. Available choices: {{'uniform'|'pseudo'|'sobol'}}.".format(
                    self.train_distribution
                )
            )

    def losses(self, targets, outputs, loss, model):
        f = []
        if self.pde is not None and get_num_args(self.pde) == 2:
            f = self.pde(model.net.inputs, outputs)
            if not isinstance(f, (list, tuple)):
                f = [f]

        def losses_train():
            f_train = f
            if self.pde is not None and get_num_args(self.pde) == 3:
                f_train = self.pde(model.net.inputs, outputs, self.train_x)
                if not isinstance(f_train, (list, tuple)):
                    f_train = [f_train]

            bcs_start = np.cumsum([0] + self.num_bcs)
            error_f = [fi[bcs_start[-1] :] for fi in f_train]
            losses = [
                loss(tf.zeros(tf.shape(error), dtype=config.real(tf)), error)
                for error in error_f
            ]
            for i, bc in enumerate(self.bcs):
                beg, end = bcs_start[i], bcs_start[i + 1]
                error = bc.error(self.train_x, model.net.inputs, outputs, beg, end)
                losses.append(
                    loss(tf.zeros(tf.shape(error), dtype=config.real(tf)), error)
                )
            return losses

        def losses_test():
            f_test = f
            if self.pde is not None and get_num_args(self.pde) == 3:
                f_test = self.pde(model.net.inputs, outputs, self.test_x)
                if not isinstance(f_test, (list, tuple)):
                    f_test = [f_test]
            return [
                loss(tf.zeros(tf.shape(fi), dtype=config.real(tf)), fi) for fi in f_test
            ] + [tf.constant(0, dtype=config.real(tf)) for _ in self.bcs]

        return tf.cond(tf.equal(model.net.data_id, 0), losses_train, losses_test)

    @run_if_all_none("train_x", "train_y")
    def train_next_batch(self, batch_size=None):
        self.train_x_all = self.train_points()
        self.train_x = self.bc_points()
        if self.pde is not None:
            self.train_x = np.vstack((self.train_x, self.train_x_all))
        self.train_y = self.soln(self.train_x) if self.soln else None
        return self.train_x, self.train_y

    @run_if_all_none("test_x", "test_y")
    def test(self):
        if self.num_test is None:
            self.test_x = self.train_x_all
        else:
            self.test_x = self.test_points()
        self.test_y = self.soln(self.test_x) if self.soln else None
        return self.test_x, self.test_y

    def resample_train_points(self):
        """Resample the training residual points.

        Warning: After resampling, need to call ``Model.compile()`` to update the loss.
        """
        self.train_x, self.train_y = None, None
        self.train_next_batch()

    def add_anchors(self, anchors):
        """Add new anchors into the training residual points.

        Warning: After adding anchors, need to call ``Model.compile()`` to update the loss.
        """
        if self.anchors is None:
            self.anchors = anchors
        else:
            self.anchors = np.vstack((anchors, self.anchors))
        self.train_x_all = np.vstack((anchors, self.train_x_all))
        self.train_x = self.bc_points()
        if self.pde is not None:
            self.train_x = np.vstack((self.train_x, self.train_x_all))
        self.train_y = self.soln(self.train_x) if self.soln else None

    def train_points(self):
        X = np.empty((0, self.geom.dim))
        if self.num_domain > 0:
            if self.train_distribution == "uniform":
                X = self.geom.uniform_points(self.num_domain, boundary=False)
            elif self.train_distribution in ["pseudo", "sobol"]:
                X = self.geom.random_points(
                    self.num_domain, random=self.train_distribution
                )
        if self.num_boundary > 0:
            if self.train_distribution == "uniform":
                tmp = self.geom.uniform_boundary_points(self.num_boundary)
            elif self.train_distribution in ["pseudo", "sobol"]:
                tmp = self.geom.random_boundary_points(
                    self.num_boundary, random=self.train_distribution
                )
            X = np.vstack((tmp, X))
        if self.anchors is not None:
            X = np.vstack((self.anchors, X))
        return X

    def bc_points(self):
        x_bcs = [bc.collocation_points(self.train_x_all) for bc in self.bcs]
        self.num_bcs = list(map(len, x_bcs))
        return np.vstack(x_bcs) if x_bcs else np.empty([0, self.train_x_all.shape[-1]])

    def test_points(self):
        return self.geom.uniform_points(self.num_test, True)


class TimePDE(PDE):
    """Time-dependent PDE solver.

    Args:
        num_domain: Number of f training points.
        num_boundary: Number of boundary condition points on the geometry boundary.
        num_initial: Number of initial condition points.
    """

    def __init__(
        self,
        geometryxtime,
        pde,
        ic_bcs,
        num_domain=0,
        num_boundary=0,
        num_initial=0,
        train_distribution="sobol",
        anchors=None,
        solution=None,
        num_test=None,
    ):
        self.num_initial = num_initial
        super(TimePDE, self).__init__(
            geometryxtime,
            pde,
            ic_bcs,
            num_domain,
            num_boundary,
            train_distribution=train_distribution,
            anchors=anchors,
            solution=solution,
            num_test=num_test,
        )

    def train_points(self):
        X = super(TimePDE, self).train_points()
        if self.num_initial > 0:
            if self.train_distribution == "uniform":
                tmp = self.geom.uniform_initial_points(self.num_initial)
            elif self.train_distribution in ["pseudo", "sobol"]:
                tmp = self.geom.random_initial_points(
                    self.num_initial, random=self.train_distribution
                )
            X = np.vstack((tmp, X))
        return X
