""" A distributed version of the paraboloid model with an extra input that can be used to shift
each index.

This version is used for testing, so it will have different options.
"""
import numpy as np

import openmdao.api as om
from openmdao.utils.mpi import MPI
from openmdao.utils.array_utils import evenly_distrib_idxs


class DistParab(om.ExplicitComponent):

    def initialize(self):
        self.options.declare('arr_size', types=int, default=10,
                             desc="Size of input and output vectors.")

        self.options.declare('deriv_type', default='dense',
                             values=['dense', 'fd', 'cs', 'sparse'],
                             desc="Method for computing derivatives.")

    def setup(self):
        arr_size = self.options['arr_size']
        deriv_type = self.options['deriv_type']
        comm = self.comm
        rank = comm.rank

        sizes, offsets = evenly_distrib_idxs(comm.size, arr_size)
        start = offsets[rank]
        io_size = sizes[rank]
        self.offset = offsets[rank]
        end = start + io_size

        self.add_input('x', val=np.ones(io_size), distributed=True,
                       src_indices=np.arange(start, end, dtype=int))
        self.add_input('y', val=np.ones(io_size), distributed=True,
                       src_indices=np.arange(start, end, dtype=int))
        self.add_input('a', val=-3.0 * np.ones(io_size), distributed=True,
                       src_indices=np.arange(start, end, dtype=int))

        self.add_output('f_xy', val=np.ones(io_size), distributed=True)

        if deriv_type == 'dense':
            self.declare_partials('f_xy', ['x', 'y', 'a'])

        elif deriv_type == 'sparse':
            row_col = np.arange(io_size)
            self.declare_partials('f_xy', ['x', 'y', 'a'], rows=row_col, cols=row_col)

        else:
            self.declare_partials('f_xy', ['x', 'y', 'a'], method=deriv_type)

    def compute(self, inputs, outputs):
        x = inputs['x']
        y = inputs['y']
        a = inputs['a']

        outputs['f_xy'] = (x + a)**2 + x * y + (y + 4.0)**2 - 3.0

    def compute_partials(self, inputs, partials):
        deriv_type = self.options['deriv_type']
        x = inputs['x']
        y = inputs['y']
        a = inputs['a']

        if deriv_type == 'dense':
            partials['f_xy', 'x'] = np.diag(2.0 * x + 2.0 * a + y)
            partials['f_xy', 'y'] = np.diag(2.0 * y + 8.0 + x)
            partials['f_xy', 'a'] = np.diag(2.0 * a + 2.0 * x)

        elif deriv_type == 'sparse':
            partials['f_xy', 'x'] = 2.0 * x + 2.0 * a + y
            partials['f_xy', 'y'] = 2.0 * y + 8.0 + x
            partials['f_xy', 'a'] = 2.0 * a + 2.0 * x


# Simplified version for feature docs without the extra testing args.

class DistParabFeature(om.ExplicitComponent):

    def initialize(self):
        self.options.declare('arr_size', types=int, default=10,
                             desc="Size of input and output vectors.")

    def setup(self):

        arr_size = self.options['arr_size']
        self.add_input('x', val=1., distributed=False,
                       shape=arr_size)
        self.add_input('y', val=1., distributed=False,
                       shape=arr_size)

        sizes, offsets = evenly_distrib_idxs(self.comm.size, arr_size)
        self.start = offsets[self.comm.rank]
        self.end = self.start + sizes[self.comm.rank]
        self.a = -3.0 + 0.6 * np.arange(self.start,self.end)

        self.add_output('f_xy', shape=len(self.a), distributed=True)
        self.add_output('f_sum', shape=1, distributed=False)

        self.declare_coloring(wrt='*', method='fd')

    def compute(self, inputs, outputs):
        x = inputs['x'][self.start:self.end]
        y = inputs['y'][self.start:self.end]

        outputs['f_xy'] = (x + self.a)**2 + x * y + (y + 4.0)**2 - 3.0

        local_sum = np.sum(outputs['f_xy'])
        global_sum = np.zeros(1)
        self.comm.Allreduce(local_sum, global_sum, op=MPI.SUM)
        outputs['f_sum'] = global_sum


class DistParabDeprecated(om.ExplicitComponent):

    def initialize(self):
        self.options.declare('arr_size', types=int, default=10,
                             desc="Size of input and output vectors.")

    def setup(self):
        arr_size = self.options['arr_size']
        comm = self.comm
        rank = comm.rank

        sizes, offsets = evenly_distrib_idxs(comm.size, arr_size)
        start = offsets[rank]
        io_size = sizes[rank]
        self.offset = offsets[rank]
        end = start + io_size

        self.add_input('x', val=np.ones(io_size), distributed=True,
                       src_indices=np.arange(start, end, dtype=int))
        self.add_input('y', val=np.ones(io_size), distributed=True,
                       src_indices=np.arange(start, end, dtype=int))
        self.add_input('offset', val=-3.0 * np.ones(io_size), distributed=True,
                       src_indices=np.arange(start, end, dtype=int))

        self.add_output('f_xy', val=np.ones(io_size), distributed=True)

        row_col = np.arange(io_size)
        self.declare_partials('f_xy', ['x', 'y', 'offset'], rows=row_col, cols=row_col)

    def compute(self, inputs, outputs):
        x = inputs['x']
        y = inputs['y']
        a = inputs['offset']

        outputs['f_xy'] = (x + a)**2 + x * y + (y + 4.0)**2 - 3.0

    def compute_partials(self, inputs, partials):
        x = inputs['x']
        y = inputs['y']
        a = inputs['offset']

        partials['f_xy', 'x'] = 2.0 * x + 2.0 * a + y
        partials['f_xy', 'y'] = 2.0 * y + 8.0 + x
        partials['f_xy', 'offset'] = 2.0 * a + 2.0 * x
