#!/usr/bin/env python
"""
This module holds classes and supporting code for simulating the vertical
plume that is generated by an underwater blowout.
"""

import six

import numpy
np = numpy

import matplotlib
from matplotlib import pylab, mlab, pyplot
plt = pyplot

from pylab import *


from datetime import datetime, timedelta

from gnome.basic_types import world_point


def get_plume_data():
    '''
      Here we will generate plume data that conforms to that of Socolofsky's
      model.
      - We will represent the mass flux amounts in kg/s.
      - We will probably get these values from a running
        plume model, but for right now, we just return
        an array with some hardcoded values.
      - For now, we return data in the format [(depth, mass_flux), ...]
    '''
    plume_mass_flux = np.zeros((10))
    plume_mass_flux[:] = 5.  # background values.
    plume_mass_flux[3] = 15.  # and now a few spikes.
    plume_mass_flux[5] = 20.
    plume_mass_flux[7] = 30.

    plume_depths = np.linspace(0, 200, plume_mass_flux.size)

    return zip(plume_depths, plume_mass_flux)


class Plume(object):
    '''
      Here, we represent the characteristics of the plume
      as a set of mass fluxes along a vertical stack.
    '''
    def __init__(self, position, plume_data):
        '''
        :param position: location the elements are released
        :type position: 3-tuple of floats (long, lat, z)
                        (note: we don't really use the z,
                         since the plume data supplies the
                         depths)

        :param plume_data: data that is supplied by a plume model
        :type plume_data: list of items in the format:
                          [(depth, mass_flux), ...]
        '''
        self.mass_flux = np.array([d[1] for d in plume_data])
        num_points = self.mass_flux.size

        self.coords = np.zeros((num_points), dtype=world_point)
        self.coords[:] = position
        self.coords[:]['z'] = np.array([d[0] for d in plume_data])


class PlumeGenerator(object):
    '''
      Here we define the method for generating LEs from a 3D plume
      over a range of time.
    '''
    def __init__(self,
                 release_time, end_release_time, time_step_delta,
                 plume):
        self.release_time = release_time
        self.end_release_time = end_release_time
        self.time_step_delta = time_step_delta

        self.plume = plume

        # Here we just calculate a reasonable value for the mass
        # that is contained in a single LE.
        # This may not be a good assumption, as other things may be
        # determining the mass of an LE in the model.
        # But we can always change this after class initialization
        # if we need to.
        self.mass_of_an_le = self.plume.mass_flux.mean() * time_step_delta

    @property
    def time_step_delta(self):
        return self._time_step_delta

    @time_step_delta.setter
    def time_step_delta(self, val):
        if type(val) not in six.integer_types + (float,) or val == 0:
            raise ValueError('time_step_delta needs to be a non-zero number')
        else:
            self._time_step_delta = val
            if self.end_release_time is not None:
                self.time_steps = (self.end_release_time - self.release_time).total_seconds()
                self.time_steps /= self._time_step_delta
            else:
                self.time_steps = None

    def _seconds_from_beginning(self, time):
        '''
          qualify our time value and return total number
          of seconds from release time to the specified time
        '''
        if time < self.release_time:
            time = self.release_time
        elif self.end_release_time is not None and time > self.end_release_time:
            time = self.end_release_time

        return (time - self.release_time).total_seconds()

    def set_le_mass_from_total_le_count(self, num_elements):
        # Calculate the mass of an LE if a total number of LEs
        # is spread across the total released mass.
        if self.end_release_time is None:
            raise OverflowError('end_release_time is undefined, '
                                'so this calculation is impossible!')
        else:
            total_release_time = (self.end_release_time - self.release_time).total_seconds()
            total_mass = self.plume.mass_flux.sum() * total_release_time
            self.mass_of_an_le = total_mass / num_elements

    def _mass_to_elems(self, mass):
        '''
          Calculate mass into an equivalent number of LEs and return them.
          - We do not count the fractional amounts.
        '''
        return (mass / self.mass_of_an_le).astype(np.int32)

    def _elems_to_mass(self, elems):
        '''
          Calculate LEs into equivalent amounts of mass and return them.
        '''
        return elems * self.mass_of_an_le

    def elems_from_beginning(self, time):
        '''
          returns the number of LEs released from release_time to the
          specified time
        '''
        seconds = self._seconds_from_beginning(time)
        mass = self.plume.mass_flux * seconds
        return self._mass_to_elems(mass)

    def elems_in_range(self, begin, end):
        return self.elems_from_beginning(end) - self.elems_from_beginning(begin)

    def __iter__(self):
        if self.time_steps is not None:
            for step in range(long(self.time_steps)):
                curr_time = self.release_time + timedelta(seconds=self.time_step_delta * step)
                next_time = curr_time + timedelta(seconds=self.time_step_delta)
                yield (curr_time,
                       zip(self.plume.coords, self.elems_in_range(curr_time,
                                                                  next_time))
                       )
        else:
            step = 0
            while True:
                curr_time = self.release_time + timedelta(seconds=self.time_step_delta * step)
                next_time = curr_time + timedelta(seconds=self.time_step_delta)
                step += 1
                yield (curr_time,
                       zip(self.plume.coords, self.elems_in_range(curr_time,
                                                                  next_time))
                       )


if __name__ == '__main__':
    release_time = datetime.now()
    end_release_time = release_time + timedelta(hours=24)
    time_step_delta = timedelta(hours=1).total_seconds()

    plume = Plume(position=(28, -78, 0.),
                  plume_data=get_plume_data())
    plume_gen = PlumeGenerator(release_time=release_time,
                               end_release_time=end_release_time,
                               time_step_delta=time_step_delta,
                               plume=plume)

    # let's print out some facts about our plume
    print '''
Based on the mean plume mass flux value,
we will choose an LE with %s kg of oil
''' % (plume_gen.mass_of_an_le)

    # now lets iterate our plume generator
    print 'First, just the occurrence pattern for LE releases...'
    total_le_count = 0
    for step in plume_gen:
        le_count = sum([r[1] for r in step[1]])
        total_le_count += le_count
        print step[0], [r[1] for r in step[1]], le_count
    print 'total LEs:', total_le_count

    print '\nNext, the full information...'
    for step in plume_gen:
        for r in step:
            # each row should consist of a world_point
            # and a number of LEs to create
            print r
        print

    # OK, we need to deal somehow with an undefined end time
    # in this case we would presumably iterate over an infinite
    # number of times.  But the rates at which we release LEs
    # should be the same.  Let's test this...
    plume_gen = PlumeGenerator(release_time=release_time,
                               end_release_time=None,  # no end in sight
                               time_step_delta=time_step_delta,
                               plume=plume)

    print '\nNext, we iterate a portion of an infinite run...'
    total_le_count = 0
    step_count = 0
    for step in plume_gen:
        le_count = sum([r[1] for r in step[1]])
        total_le_count += le_count
        print step[0], [r[1] for r in step[1]], le_count

        step_count += 1
        if step_count >= 24:
            break
    print 'total LEs:', total_le_count

    # I believe with our test data that the total LEs is 240
    # Lets change the total number of LEs and see how that affects our run.
    plume_gen = PlumeGenerator(release_time=release_time,
                               end_release_time=end_release_time,
                               time_step_delta=time_step_delta,
                               plume=plume)
    plume_gen.set_le_mass_from_total_le_count(200)
    print '\nNow, the occurrence pattern if the total LEs is 200...'
    total_le_count = 0
    for step in plume_gen:
        le_count = sum([r[1] for r in step[1]])
        total_le_count += le_count
        print step[0], [r[1] for r in step[1]], le_count
    print 'total LEs:', total_le_count
    assert total_le_count == 200

    #
    # OK, that worked for 200 LEs, but there will be cases where an LE
    # did not show up in the output of a plume generation run.
    # I think we need to analyze how the actual number of LEs
    # might differ from the number of LEs that we are trying to specify.
    # And we need to run this test on some real plume data when we get it.
    #
    def compare_le_count(plume_generator, le_count):
        plume_generator.set_le_mass_from_total_le_count(le_count)
        return le_count, sum([sum([r[1] for r in step[1]]) for step in plume_generator])

    # To start with, we will compare the number of LEs we specified vs. the number
    # of LEs that we came up with after a run of our plume generator.  And we will
    # do this over a range of specified LE counts.
    le_counts = [compare_le_count(plume_gen, n) for n in range(100, 401)]
    print le_counts

    figure(num=10)
    plot([n[0] for n in le_counts], le_counts)
    title('Specified LEs vs. Resulting LEs')

    # Here we plot the difference in LEs that we had.
    # - I would expect there is a small chance where
    #   every one of our points missed an LE.  In this case
    #   we would be off by an amount equal to the number of points in our
    #   plume column.
    # - I would not expect any single point to miss more than one LE.
    #   Thus, I would not expect a number greater than the number of
    #   plume data points.
    # - For our test data, this maximum number is 10, and it
    #   occurs when we specify 260 LEs.
    figure(num=20)
    plot([n[0] for n in le_counts], [diff(i) for i in le_counts])
    xlabel('Number of LEs specified')
    ylabel('Diff in LEs')
    title('Specified LEs vs. Resulting LEs)')
    assert max([abs(diff(i)) for i in le_counts])[0] <= plume_gen.plume.mass_flux.size

    # OK the biggest deviation was about 8% when we specified 119 LEs, and
    # it tends to get smaller as we increase the specified LE count.
    figure(num=30)
    plot([n[0] for n in le_counts],
         [(abs(diff(i)) / np.float(i[0]) * 100.) for i in le_counts])
    xlabel('Number of LEs specified')
    ylabel('Deviation (%)')
    title('Specified LEs vs. Resulting LEs')

    show()
    #plt.close('all')
