# (C) British Crown Copyright 2010 - 2013, Met Office
#
# This file is part of Iris.
#
# Iris is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the
# Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Iris is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Iris.  If not, see <http://www.gnu.org/licenses/>.

# import iris tests first so that some things can be initialised before
# importing anything else
import iris.tests as tests

from functools import wraps
import types
import warnings

import matplotlib.pyplot as plt
import numpy as np

import iris
import iris.coords as coords
import iris.plot as iplt
import iris.quickplot as qplt
import iris.symbols
import iris.tests.stock
import iris.tests.test_mapping as test_mapping

def simple_cube():
    cube = iris.tests.stock.realistic_4d()
    cube = cube[:, 0, 0, :]
    cube.coord('time').guess_bounds()
    return cube

class TestSimple(tests.GraphicsTest):
    def test_points(self):
        cube = simple_cube()
        qplt.contourf(cube)
        self.check_graphic()

    def test_bounds(self):
        cube = simple_cube()
        qplt.pcolor(cube)
        self.check_graphic()

class TestMissingCoord(tests.GraphicsTest):
    def _check(self, cube):
        qplt.contourf(cube)
        self.check_graphic()

        qplt.pcolor(cube)
        self.check_graphic()

    def test_no_u(self):
        cube = simple_cube()
        cube.remove_coord('grid_longitude')
        self._check(cube)

    def test_no_v(self):
        cube = simple_cube()
        cube.remove_coord('time')
        self._check(cube)

    def test_none(self):
        cube = simple_cube()
        cube.remove_coord('grid_longitude')
        cube.remove_coord('time')
        self._check(cube)

@iris.tests.skip_data
class TestMissingCS(tests.GraphicsTest):
    @iris.tests.skip_data
    def test_missing_cs(self):
        cube = tests.stock.simple_pp()
        cube.coord("latitude").coord_system = None
        cube.coord("longitude").coord_system = None
        qplt.contourf(cube)
        qplt.plt.gca().coastlines()
        self.check_graphic()

class TestHybridHeight(tests.GraphicsTest):
    def setUp(self):
        self.cube = iris.tests.stock.realistic_4d()[0, :15, 0, :]

    def _check(self, plt_method, test_altitude=True):
        plt_method(self.cube)
        self.check_graphic()

        plt_method(self.cube, coords=['level_height', 'grid_longitude'])
        self.check_graphic()

        plt_method(self.cube, coords=['grid_longitude', 'level_height'])
        self.check_graphic()

        if test_altitude:
            plt_method(self.cube, coords=['grid_longitude', 'altitude'])
            self.check_graphic()

            plt_method(self.cube, coords=['altitude', 'grid_longitude'])
            self.check_graphic()

    def test_points(self):
        self._check(qplt.contourf)

    def test_bounds(self):
        self._check(qplt.pcolor, test_altitude=False)

    def test_orography(self):
        qplt.contourf(self.cube)
        iplt.orography_at_points(self.cube)
        iplt.points(self.cube)
        self.check_graphic()

        coords = ['altitude', 'grid_longitude']
        qplt.contourf(self.cube, coords=coords)
        iplt.orography_at_points(self.cube, coords=coords)
        iplt.points(self.cube, coords=coords)
        self.check_graphic()

        # TODO: Test bounds once they are supported.
        with self.assertRaises(NotImplementedError):
            qplt.pcolor(self.cube)
            iplt.orography_at_bounds(self.cube)
            iplt.outline(self.cube)
            self.check_graphic()

class Test1dPlotMultiArgs(tests.GraphicsTest):
    # tests for iris.plot using multi-argument calling convention

    def setUp(self):
        self.cube1d = _load_4d_testcube()[0, :, 0, 0]
        self.draw_method = iplt.plot

    def test_cube(self):
        # just plot a cube against its dim coord
        self.draw_method(self.cube1d)   # altitude vs temp
        self.check_graphic()

    def test_coord(self):
        # plot the altitude coordinate
        self.draw_method(self.cube1d.coord('altitude'))
        self.check_graphic()

    def test_coord_cube(self):
        # plot temperature against sigma
        self.draw_method(self.cube1d.coord('sigma'), self.cube1d)
        self.check_graphic()

    def test_cube_coord(self):
        # plot a vertical profile of temperature
        self.draw_method(self.cube1d, self.cube1d.coord('altitude'))
        self.check_graphic()

    def test_coord_coord(self):
        # plot two coordinates that are not mappable
        self.draw_method(self.cube1d.coord('sigma'),
                         self.cube1d.coord('altitude'))
        self.check_graphic()

    def test_coord_coord_map(self):
        # plot lat-lon aux coordinates of a trajectory, which draws a map
        lon = iris.coords.AuxCoord([0, 5, 10, 15, 20, 25, 30, 35, 40, 45],
                                   standard_name='longitude',
                                   units='degrees_north')
        lat = iris.coords.AuxCoord([45, 55, 50, 60, 55, 65, 60, 70, 65, 75],
                                   standard_name='latitude',
                                   units='degrees_north')
        self.draw_method(lon, lat)
        plt.gca().coastlines()
        self.check_graphic()

    def test_cube_cube(self):
        # plot two phenomena against eachother, in this case just dummy data
        cube1 = self.cube1d.copy()
        cube2 = self.cube1d.copy()
        cube1.rename('some phenomenon')
        cube2.rename('some other phenomenon')
        cube1.units = iris.unit.Unit('no_unit')
        cube2.units = iris.unit.Unit('no_unit')
        cube1.data[:] = np.linspace(0, 1, 7)
        cube2.data[:] = np.exp(cube1.data)
        self.draw_method(cube1, cube2)
        self.check_graphic()

    def test_incompatible_objects(self):
        # incompatible objects (not the same length) should raise an error
        with self.assertRaises(ValueError):
            self.draw_method(self.cube1d.coord('time'), (self.cube1d))

    def test_multimidmensional(self):
        # multidimensional cubes are not allowed
        cube = _load_4d_testcube()[0, :, :, 0]
        with self.assertRaises(ValueError):
            self.draw_method(cube)

    def test_not_cube_or_coord(self):
        # inputs must be cubes or coordinates, otherwise an error should be
        # raised
        xdim = np.arange(self.cube1d.shape[0])
        with self.assertRaises(TypeError):
            self.draw_method(xdim, self.cube1d)

    def test_coords_deprecated(self):
        # ensure a warning is raised if the old coords keyword argument is
        # used, and make sure the plot produced is consistent with the old
        # interface
        msg = 'Missing deprecation warning for coords keyword.'
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always')
            self.draw_method(self.cube1d, coords=['sigma'])
            self.assertEqual(len(w), 1, msg)
        self.check_graphic()

    def test_coords_deprecation_too_many(self):
        # in deprecation mode, too many coords is an error
        with self.assertRaises(ValueError):
            self.draw_method(self.cube1d, coords=['sigma', 'sigma'])

    def test_coords_deprecation_invalid_span(self):
        # in deprecation mode, a coordinate that doesn't span data is an error
        with self.assertRaises(ValueError):
            self.draw_method(self.cube1d, coords=['time'])

class Test1dQuickplotPlotMultiArgs(Test1dPlotMultiArgs):
    # tests for iris.plot using multi-argument calling convention

    def setUp(self):
        self.cube1d = _load_4d_testcube()[0, :, 0, 0]
        self.draw_method = qplt.plot

@tests.skip_data
class Test1dScatter(tests.GraphicsTest):

    def setUp(self):
        self.cube = iris.load_cube(
            tests.get_data_path(('NAME', 'NAMEIII_trajectory.txt')),
            'Temperature')
        self.draw_method = iplt.scatter

    def test_coord_coord(self):
        x = self.cube.coord('longitude')
        y = self.cube.coord('height')
        c = self.cube.data
        self.draw_method(x, y, c=c, edgecolor='none')
        self.check_graphic()

    def test_coord_coord_map(self):
        x = self.cube.coord('longitude')
        y = self.cube.coord('latitude')
        c = self.cube.data
        self.draw_method(x, y, c=c, edgecolor='none')
        plt.gca().coastlines()
        self.check_graphic()

    def test_coord_cube(self):
        x = self.cube.coord('latitude')
        y = self.cube
        c = self.cube.coord('Travel Time').points
        self.draw_method(x, y, c=c, edgecolor='none')
        self.check_graphic()

    def test_cube_coord(self):
        x = self.cube
        y = self.cube.coord('height')
        c = self.cube.coord('Travel Time').points
        self.draw_method(x, y, c=c, edgecolor='none')
        self.check_graphic()

    def test_cube_cube(self):
        x = iris.load_cube(
            tests.get_data_path(('NAME', 'NAMEIII_trajectory.txt')),
            'Rel Humidity')
        y = self.cube
        c = self.cube.coord('Travel Time').points
        self.draw_method(x, y, c=c, edgecolor='none')
        self.check_graphic()

    def test_incompatible_objects(self):
        # cubes/coordinates of different sizes cannot be plotted
        x = self.cube
        y = self.cube.coord('height')[:-1]
        with self.assertRaises(ValueError):
            self.draw_method(x, y)

    def test_multidimensional(self):
        # multidimensional cubes/coordinates are not allowed
        x = _load_4d_testcube()[0, :, :, 0]
        y = x.coord('model_level_number')
        with self.assertRaises(ValueError):
            self.draw_method(x, y)

    def test_not_cube_or_coord(self):
        # inputs must be cubes or coordinates
        x = np.arange(self.cube.shape[0])
        y = self.cube
        with self.assertRaises(TypeError):
            self.draw_method(x, y)

@tests.skip_data
class Test1dQuickplotScatter(Test1dScatter):

    def setUp(self):
        self.cube = iris.load_cube(
            tests.get_data_path(('NAME', 'NAMEIII_trajectory.txt')),
            'Temperature')
        self.draw_method = qplt.scatter

@iris.tests.skip_data
class TestAttributePositive(tests.GraphicsTest):
    def test_1d_positive_up(self):
        path = tests.get_data_path(('NetCDF', 'ORCA2', 'votemper.nc'))
        cube = iris.load_cube(path)
        qplt.plot(cube.coord('depth'), cube[0, :, 60, 80])
        self.check_graphic()

    def test_1d_positive_down(self):
        path = tests.get_data_path(('NetCDF', 'ORCA2', 'votemper.nc'))
        cube = iris.load_cube(path)
        qplt.plot(cube[0, :, 60, 80], cube.coord('depth'))
        self.check_graphic()

    def test_2d_positive_up(self):
        path = tests.get_data_path(('NetCDF', 'testing',
                                    'small_theta_colpex.nc'))
        cube = iris.load_cube(path)[0, :, 42, :]
        qplt.pcolormesh(cube)
        self.check_graphic()

    def test_2d_positive_down(self):
        path = tests.get_data_path(('NetCDF', 'ORCA2', 'votemper.nc'))
        cube = iris.load_cube(path)[0, :, 42, :]
        qplt.pcolormesh(cube)
        self.check_graphic()

# Caches _load_4d_testcube so subsequent calls are faster
def cache(fn, cache={}):
    def inner(*args, **kwargs):
        key = fn.__name__
        if key not in cache:
            cache[key] = fn(*args, **kwargs)
        return cache[key]
    return inner

@cache
def _load_4d_testcube():
    # Load example 4d data (TZYX).
    test_cube = iris.tests.stock.realistic_4d()
    # Replace forecast_period coord with a multi-valued version.
    time_coord = test_cube.coord('time')
    n_times = len(time_coord.points)
    forecast_dims = test_cube.coord_dims(time_coord)
    test_cube.remove_coord('forecast_period')
    # Make up values (including bounds), to roughly match older testdata.
    point_values = np.linspace((1 + 1.0 / 6), 2.0, n_times)
    point_uppers = point_values + (point_values[1] - point_values[0])
    bound_values = np.column_stack([point_values, point_uppers])
    # NOTE: this must be a DimCoord
    #  - an equivalent AuxCoord produces different plots.
    new_forecast_coord = iris.coords.DimCoord(
        points=point_values,
        bounds=bound_values,
        standard_name='forecast_period',
        units=iris.unit.Unit('hours')
    )
    test_cube.add_aux_coord(new_forecast_coord, forecast_dims)
    # Heavily reduce dimensions for faster testing.
    # NOTE: this makes ZYX non-contiguous.  Doesn't seem to matter for now.
    test_cube = test_cube[:, ::10, ::10, ::10]
    return test_cube

@cache
def _load_wind_no_bounds():
    # Load the COLPEX data => TZYX
    path = tests.get_data_path(('PP', 'COLPEX', 'small_eastward_wind.pp'))
    wind = iris.load_cube(path, 'eastward_wind')

    # Remove bounds from all coords that have them.
    wind.coord('grid_latitude').bounds = None
    wind.coord('grid_longitude').bounds = None
    wind.coord('level_height').bounds = None
    wind.coord('sigma').bounds = None

    return wind[:, :, :50, :50]

def _time_series(src_cube):
    # Until we have plotting support for multiple axes on the same dimension,
    # remove the time coordinate and its axis.
    cube = src_cube.copy()
    cube.remove_coord('time')
    return cube

def _date_series(src_cube):
    # Until we have plotting support for multiple axes on the same dimension,
    # remove the forecast_period coordinate and its axis.
    cube = src_cube.copy()
    cube.remove_coord('forecast_period')
    return cube

class SliceMixin(object):
    """Mixin class providing tests for each 2-dimensional permutation of axes.

    Requires self.draw_method to be the relevant plotting function,
    and self.results to be a dictionary containing the desired test results."""

    def test_yx(self):
        cube = self.wind[0, 0, :, :]
        self.draw_method(cube)
        self.check_graphic()

    def test_zx(self):
        cube = self.wind[0, :, 0, :]
        self.draw_method(cube)
        self.check_graphic()

    def test_tx(self):
        cube = _time_series(self.wind[:, 0, 0, :])
        self.draw_method(cube)
        self.check_graphic()

    def test_zy(self):
        cube = self.wind[0, :, :, 0]
        self.draw_method(cube)
        self.check_graphic()

    def test_ty(self):
        cube = _time_series(self.wind[:, 0, :, 0])
        self.draw_method(cube)
        self.check_graphic()

    def test_tz(self):
        cube = _time_series(self.wind[:, :, 0, 0])
        self.draw_method(cube)
        self.check_graphic()

@iris.tests.skip_data
class TestContour(tests.GraphicsTest, SliceMixin):
    """Test the iris.plot.contour routine."""
    def setUp(self):
        self.wind = _load_4d_testcube()
        self.draw_method = iplt.contour

@iris.tests.skip_data
class TestContourf(tests.GraphicsTest, SliceMixin):
    """Test the iris.plot.contourf routine."""
    def setUp(self):
        self.wind = _load_4d_testcube()
        self.draw_method = iplt.contourf

@iris.tests.skip_data
class TestPcolor(tests.GraphicsTest, SliceMixin):
    """Test the iris.plot.pcolor routine."""
    def setUp(self):
        self.wind = _load_4d_testcube()
        self.draw_method = iplt.pcolor

@iris.tests.skip_data
class TestPcolormesh(tests.GraphicsTest, SliceMixin):
    """Test the iris.plot.pcolormesh routine."""
    def setUp(self):
        self.wind = _load_4d_testcube()
        self.draw_method = iplt.pcolormesh

def check_warnings(method):
    """
    Decorator that adds a catch_warnings and filter to assert
    the method being decorated issues a UserWarning.

    """
    @wraps(method)
    def decorated_method(self, *args, **kwargs):
        # Force reset of iris.coords warnings registry to avoid suppression of
        # repeated warnings. warnings.resetwarnings() does not do this.
        if hasattr(coords, '__warningregistry__'):
            coords.__warningregistry__.clear()

        # Check that method raises warning.
        with warnings.catch_warnings():
            warnings.simplefilter("error")
            with self.assertRaises(UserWarning):
                return method(self, *args, **kwargs)
    return decorated_method

def ignore_warnings(method):
    """
    Decorator that adds a catch_warnings and filter to suppress
    any warnings issues by the method being decorated.

    """
    @wraps(method)
    def decorated_method(self, *args, **kwargs):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            return method(self, *args, **kwargs)
    return decorated_method

class CheckForWarningsMetaclass(type):
    """
    Metaclass that adds a further test for each base class test
    that checks that each test raises a UserWarning. Each base
    class test is then overriden to ignore warnings in order to
    check the underlying functionality.

    """
    def __new__(cls, name, bases, local):
        def add_decorated_methods(attr_dict, target_dict, decorator):
            for key, value in attr_dict.items():
                if (isinstance(value, types.FunctionType) and
                        key.startswith('test')):
                    new_key = '_'.join((key, decorator.__name__))
                    if new_key not in target_dict:
                        wrapped = decorator(value)
                        wrapped.__name__ = new_key
                        target_dict[new_key] = wrapped
                    else:
                        raise RuntimeError('A attribute called {!r} '
                                           'already exists.'.format(new_key))

        def override_with_decorated_methods(attr_dict, target_dict,
                                            decorator):
            for key, value in attr_dict.items():
                if (isinstance(value, types.FunctionType) and
                        key.startswith('test')):
                    target_dict[key] = decorator(value)

        # Add decorated versions of base methods
        # to check for warnings.
        for base in bases:
            add_decorated_methods(base.__dict__, local, check_warnings)

        # Override base methods to ignore warnings.
        for base in bases:
            override_with_decorated_methods(base.__dict__, local,
                                            ignore_warnings)

        return type.__new__(cls, name, bases, local)

@iris.tests.skip_data
class TestPcolorNoBounds(tests.GraphicsTest, SliceMixin):
    """
    Test the iris.plot.pcolor routine on a cube with coordinates
    that have no bounds.

    """
    __metaclass__ = CheckForWarningsMetaclass

    def setUp(self):
        self.wind = _load_wind_no_bounds()
        self.draw_method = iplt.pcolor

@iris.tests.skip_data
class TestPcolormeshNoBounds(tests.GraphicsTest, SliceMixin):
    """
    Test the iris.plot.pcolormesh routine on a cube with coordinates
    that have no bounds.

    """
    __metaclass__ = CheckForWarningsMetaclass

    def setUp(self):
        self.wind = _load_wind_no_bounds()
        self.draw_method = iplt.pcolormesh

class Slice1dMixin(object):
    """Mixin class providing tests for each 1-dimensional permutation of axes.

    Requires self.draw_method to be the relevant plotting function,
    and self.results to be a dictionary containing the desired test results."""

    def test_x(self):
        cube = self.wind[0, 0, 0, :]
        self.draw_method(cube)
        self.check_graphic()

    def test_y(self):
        cube = self.wind[0, 0, :, 0]
        self.draw_method(cube)
        self.check_graphic()

    def test_z(self):
        cube = self.wind[0, :, 0, 0]
        self.draw_method(cube)
        self.check_graphic()

    def test_t(self):
        cube = _time_series(self.wind[:, 0, 0, 0])
        self.draw_method(cube)
        self.check_graphic()

    def test_t_dates(self):
        cube = _date_series(self.wind[:, 0, 0, 0])
        self.draw_method(cube)
        plt.gcf().autofmt_xdate()
        plt.xlabel('Phenomenon time')

        self.check_graphic()

@iris.tests.skip_data
class TestPlot(tests.GraphicsTest, Slice1dMixin):
    """Test the iris.plot.plot routine."""
    def setUp(self):
        self.wind = _load_4d_testcube()
        self.draw_method = iplt.plot

@iris.tests.skip_data
class TestQuickplotPlot(tests.GraphicsTest, Slice1dMixin):
    """Test the iris.quickplot.plot routine."""
    def setUp(self):
        self.wind = _load_4d_testcube()
        self.draw_method = qplt.plot

_load_cube_once_cache = {}

def load_cube_once(filename, constraint):
    """Same syntax as load_cube, but will only load a file once,

    then cache the answer in a dictionary.

    """
    global _load_cube_once_cache
    key = (filename, str(constraint))
    cube = _load_cube_once_cache.get(key, None)

    if cube is None:
        cube = iris.load_cube(filename, constraint)
        _load_cube_once_cache[key] = cube

    return cube

class LambdaStr(object):
    """Provides a callable function which has a sensible __repr__."""
    def __init__(self, repr, lambda_fn):
        self.repr = repr
        self.lambda_fn = lambda_fn

    def __call__(self, *args, **kwargs):
        return self.lambda_fn(*args, **kwargs)

    def __repr__(self):
        return self.repr

@iris.tests.skip_data
class TestPlotCoordinatesGiven(tests.GraphicsTest):
    def setUp(self):
        filename = tests.get_data_path(('PP', 'COLPEX',
                                        'theta_and_orog_subset.pp'))
        self.cube = load_cube_once(filename, 'air_potential_temperature')

        self.draw_module = iris.plot
        self.contourf = LambdaStr('iris.plot.contourf',
                                  lambda cube, *args, **kwargs:
                                  iris.plot.contourf(cube, *args, **kwargs))
        self.contour = LambdaStr('iris.plot.contour',
                                 lambda cube, *args, **kwargs:
                                 iris.plot.contour(cube, *args, **kwargs))
        self.points = LambdaStr('iris.plot.points',
                                lambda cube, *args, **kwargs:
                                iris.plot.points(cube, c=cube.data,
                                                 *args, **kwargs))
        self.plot = LambdaStr('iris.plot.plot',
                              lambda cube, *args, **kwargs:
                              iris.plot.plot(cube, *args, **kwargs))

        self.results = {'yx': ([self.contourf, ['grid_latitude',
                                                'grid_longitude']],
                               [self.contourf, ['grid_longitude',
                                                'grid_latitude']],
                               [self.contour, ['grid_latitude',
                                               'grid_longitude']],
                               [self.contour, ['grid_longitude',
                                               'grid_latitude']],
                               [self.points, ['grid_latitude',
                                              'grid_longitude']],
                               [self.points, ['grid_longitude',
                                              'grid_latitude']],),
                        'zx': ([self.contourf, ['model_level_number',
                                                'grid_longitude']],
                               [self.contourf, ['grid_longitude',
                                                'model_level_number']],
                               [self.contour, ['model_level_number',
                                               'grid_longitude']],
                               [self.contour, ['grid_longitude',
                                               'model_level_number']],
                               [self.points, ['model_level_number',
                                              'grid_longitude']],
                               [self.points, ['grid_longitude',
                                              'model_level_number']],),
                        'tx': ([self.contourf, ['time', 'grid_longitude']],
                               [self.contourf, ['grid_longitude', 'time']],
                               [self.contour, ['time', 'grid_longitude']],
                               [self.contour, ['grid_longitude', 'time']],
                               [self.points, ['time', 'grid_longitude']],
                               [self.points, ['grid_longitude', 'time']],),
                        'x': ([self.plot, ['grid_longitude']],),
                        'y': ([self.plot, ['grid_latitude']],)
                        }

    def draw(self, draw_method, *args, **kwargs):
        draw_fn = getattr(self.draw_module, draw_method)
        draw_fn(*args, **kwargs)
        self.check_graphic()

    def run_tests(self, cube, results):
        for draw_method, coords in results:
            draw_method(cube, coords=coords)
            try:
                self.check_graphic()
            except AssertionError, err:
                self.fail('Draw method %r failed with coords: %r. '
                          'Assertion message: %s' % (draw_method, coords, err))

    def run_tests_1d(self, cube, results):
        # there is a different calling convention for 1d plots
        for draw_method, coords in results:
            draw_method(cube.coord(coords[0]), cube)
            try:
                self.check_graphic()
            except AssertionError as err:
                msg = 'Draw method {!r} failed with coords: {!r}. ' \
                      'Assertion message: {!s}'
                self.fail(msg.format(draw_method, coords, err))

    def test_yx(self):
        test_cube = self.cube[0, 0, :, :]
        self.run_tests(test_cube, self.results['yx'])

    def test_zx(self):
        test_cube = self.cube[0, :15, 0, :]
        self.run_tests(test_cube, self.results['zx'])

    def test_tx(self):
        test_cube = self.cube[:, 0, 0, :]
        self.run_tests(test_cube, self.results['tx'])

    def test_x(self):
        test_cube = self.cube[0, 0, 0, :]
        self.run_tests_1d(test_cube, self.results['x'])

    def test_y(self):
        test_cube = self.cube[0, 0, :, 0]
        self.run_tests_1d(test_cube, self.results['y'])

    def test_badcoords(self):
        cube = self.cube[0, 0, :, :]
        draw_fn = getattr(self.draw_module, 'contourf')
        self.assertRaises(ValueError, draw_fn, cube,
                          coords=['grid_longitude', 'grid_longitude'])
        self.assertRaises(ValueError, draw_fn, cube,
                          coords=['grid_longitude', 'grid_longitude',
                                  'grid_latitude'])
        self.assertRaises(iris.exceptions.CoordinateNotFoundError, draw_fn,
                          cube, coords=['grid_longitude', 'wibble'])
        self.assertRaises(ValueError, draw_fn, cube, coords=[])
        self.assertRaises(ValueError, draw_fn, cube,
                          coords=[cube.coord('grid_longitude'),
                                  cube.coord('grid_longitude')])
        self.assertRaises(ValueError, draw_fn, cube,
                          coords=[cube.coord('grid_longitude'),
                                  cube.coord('grid_longitude'),
                                  cube.coord('grid_longitude')])

    def test_non_cube_coordinate(self):
        cube = self.cube[0, :, :, 0]
        pts = -100 + np.arange(cube.shape[1]) * 13
        x = coords.DimCoord(pts, standard_name='model_level_number',
                            attributes={'positive': 'up'})
        self.draw('contourf', cube, coords=['grid_latitude', x])

@iris.tests.skip_data
class TestPlotDimAndAuxCoordsKwarg(tests.GraphicsTest):
    def setUp(self):
        filename = tests.get_data_path(('NetCDF', 'rotated', 'xy',
                                        'rotPole_landAreaFraction.nc'))
        self.cube = iris.load_cube(filename)

    def test_default(self):
        iplt.contourf(self.cube)
        plt.gca().coastlines()
        self.check_graphic()

    def test_coords(self):
        # Pass in dimension coords.
        rlat = self.cube.coord('grid_latitude')
        rlon = self.cube.coord('grid_longitude')
        iplt.contourf(self.cube, coords=[rlon, rlat])
        plt.gca().coastlines()
        self.check_graphic()
        # Pass in auxiliary coords.
        lat = self.cube.coord('latitude')
        lon = self.cube.coord('longitude')
        iplt.contourf(self.cube, coords=[lon, lat])
        plt.gca().coastlines()
        self.check_graphic()

    def test_coord_names(self):
        # Pass in names of dimension coords.
        iplt.contourf(self.cube, coords=['grid_longitude', 'grid_latitude'])
        plt.gca().coastlines()
        self.check_graphic()
        # Pass in names of auxiliary coords.
        iplt.contourf(self.cube, coords=['longitude', 'latitude'])
        plt.gca().coastlines()
        self.check_graphic()

    def test_yx_order(self):
        # Do not attempt to draw coastlines as it is not a map.
        iplt.contourf(self.cube, coords=['grid_latitude', 'grid_longitude'])
        self.check_graphic()
        iplt.contourf(self.cube, coords=['latitude', 'longitude'])
        self.check_graphic()

class TestSymbols(tests.GraphicsTest):
    def test_cloud_cover(self):
        iplt.symbols(range(10), [0] * 10, [iris.symbols.CLOUD_COVER[i]
                                           for i in range(10)], 0.375)
        self.check_graphic()

class TestPlottingExceptions(tests.IrisTest):
    def setUp(self):
        self.bounded_cube = tests.stock.lat_lon_cube()
        self.bounded_cube.coord("latitude").guess_bounds()
        self.bounded_cube.coord("longitude").guess_bounds()

    def test_boundmode_multidim(self):
        # Test exception translation.
        # We can't get contiguous bounded grids from multi-d coords.
        cube = self.bounded_cube
        cube.remove_coord("latitude")
        cube.add_aux_coord(coords.AuxCoord(points=cube.data,
                                           standard_name='latitude',
                                           units='degrees'), [0, 1])
        with self.assertRaises(ValueError):
            iplt.pcolormesh(cube, coords=['longitude', 'latitude'])

    def test_boundmode_4bounds(self):
        # Test exception translation.
        # We can only get contiguous bounded grids with 2 bounds per point.
        cube = self.bounded_cube
        lat = coords.AuxCoord.from_coord(cube.coord("latitude"))
        lat.bounds = np.array([lat.points, lat.points + 1,
                               lat.points + 2, lat.points + 3]).transpose()
        cube.remove_coord("latitude")
        cube.add_aux_coord(lat, 0)
        with self.assertRaises(ValueError):
            iplt.pcolormesh(cube, coords=['longitude', 'latitude'])

    def test_different_coord_systems(self):
        cube = self.bounded_cube
        lat = cube.coord('latitude')
        lon = cube.coord('longitude')
        lat.coord_system = iris.coord_systems.GeogCS(7000000)
        lon.coord_system = iris.coord_systems.GeogCS(7000001)
        with self.assertRaises(ValueError):
            iplt.pcolormesh(cube, coords=['longitude', 'latitude'])

@iris.tests.skip_data
class TestPlotOtherCoordSystems(tests.GraphicsTest):
    def test_plot_tmerc(self):
        filename = tests.get_data_path(('NetCDF', 'transverse_mercator',
                                        'tmean_1910_1910.nc'))
        self.cube = iris.load_cube(filename)
        iplt.pcolormesh(self.cube[0])
        plt.gca().coastlines()
        self.check_graphic()

if __name__ == "__main__":
    tests.main()


# Author: Arnaud Joly, Joel Nothman, Hamzeh Alsalhi
#
# License: BSD 3 clause
"""
Multi-class / multi-label utility function
==========================================

"""
from __future__ import division
from collections import Sequence
from itertools import chain

from scipy.sparse import issparse
from scipy.sparse.base import spmatrix
from scipy.sparse import dok_matrix
from scipy.sparse import lil_matrix

import numpy as np

from ..externals.six import string_types
from .validation import check_array
from ..utils.fixes import bincount
from ..utils.fixes import array_equal

def _unique_multiclass(y):
    if hasattr(y, '__array__'):
        return np.unique(np.asarray(y))
    else:
        return set(y)

def _unique_indicator(y):
    return np.arange(check_array(y, ['csr', 'csc', 'coo']).shape[1])

_FN_UNIQUE_LABELS = {
    'binary': _unique_multiclass,
    'multiclass': _unique_multiclass,
    'multilabel-indicator': _unique_indicator,
}

def unique_labels(*ys):
    """Extract an ordered array of unique labels

    We don't allow:
        - mix of multilabel and multiclass (single label) targets
        - mix of label indicator matrix and anything else,
          because there are no explicit labels)
        - mix of label indicator matrices of different sizes
        - mix of string and integer labels

    At the moment, we also don't allow "multiclass-multioutput" input type.

    Parameters
    ----------
    *ys : array-likes,

    Returns
    -------
    out : numpy array of shape [n_unique_labels]
        An ordered array of unique labels.

    Examples
    --------
    >>> from sklearn.utils.multiclass import unique_labels
    >>> unique_labels([3, 5, 5, 5, 7, 7])
    array([3, 5, 7])
    >>> unique_labels([1, 2, 3, 4], [2, 2, 3, 4])
    array([1, 2, 3, 4])
    >>> unique_labels([1, 2, 10], [5, 11])
    array([ 1,  2,  5, 10, 11])
    """
    if not ys:
        raise ValueError('No argument has been passed.')
    # Check that we don't mix label format

    ys_types = set(type_of_target(x) for x in ys)
    if ys_types == set(["binary", "multiclass"]):
        ys_types = set(["multiclass"])

    if len(ys_types) > 1:
        raise ValueError("Mix type of y not allowed, got types %s" % ys_types)

    label_type = ys_types.pop()

    # Check consistency for the indicator format
    if (label_type == "multilabel-indicator" and
            len(set(check_array(y, ['csr', 'csc', 'coo']).shape[1]
                    for y in ys)) > 1):
        raise ValueError("Multi-label binary indicator input with "
                         "different numbers of labels")

    # Get the unique set of labels
    _unique_labels = _FN_UNIQUE_LABELS.get(label_type, None)
    if not _unique_labels:
        raise ValueError("Unknown label type: %s" % repr(ys))

    ys_labels = set(chain.from_iterable(_unique_labels(y) for y in ys))

    # Check that we don't mix string type with number type
    if (len(set(isinstance(label, string_types) for label in ys_labels)) > 1):
        raise ValueError("Mix of label input types (string and number)")

    return np.array(sorted(ys_labels))

def _is_integral_float(y):
    return y.dtype.kind == 'f' and np.all(y.astype(int) == y)

def is_multilabel(y):
    """ Check if ``y`` is in a multilabel format.

    Parameters
    ----------
    y : numpy array of shape [n_samples]
        Target values.

    Returns
    -------
    out : bool,
        Return ``True``, if ``y`` is in a multilabel format, else ```False``.

    Examples
    --------
    >>> import numpy as np
    >>> from sklearn.utils.multiclass import is_multilabel
    >>> is_multilabel([0, 1, 0, 1])
    False
    >>> is_multilabel([[1], [0, 2], []])
    False
    >>> is_multilabel(np.array([[1, 0], [0, 0]]))
    True
    >>> is_multilabel(np.array([[1], [0], [0]]))
    False
    >>> is_multilabel(np.array([[1, 0, 0]]))
    True
    """
    if hasattr(y, '__array__'):
        y = np.asarray(y)
    if not (hasattr(y, "shape") and y.ndim == 2 and y.shape[1] > 1):
        return False

    if issparse(y):
        if isinstance(y, (dok_matrix, lil_matrix)):
            y = y.tocsr()
        return (len(y.data) == 0 or np.unique(y.data).size == 1 and
                (y.dtype.kind in 'biu' or  # bool, int, uint
                 _is_integral_float(np.unique(y.data))))
    else:
        labels = np.unique(y)

        return len(labels) < 3 and (y.dtype.kind in 'biu' or  # bool, int, uint
                                    _is_integral_float(labels))

def check_classification_targets(y):
    """Ensure that target y is of a non-regression type.

    Only the following target types (as defined in type_of_target) are allowed:
        'binary', 'multiclass', 'multiclass-multioutput',
        'multilabel-indicator', 'multilabel-sequences'

    Parameters
    ----------
    y : array-like
    """
    y_type = type_of_target(y)
    if y_type not in ['binary', 'multiclass', 'multiclass-multioutput',
            'multilabel-indicator', 'multilabel-sequences']:
        raise ValueError("Unknown label type: %r" % y_type)

def type_of_target(y):
    """Determine the type of data indicated by target `y`

    Parameters
    ----------
    y : array-like

    Returns
    -------
    target_type : string
        One of:
        * 'continuous': `y` is an array-like of floats that are not all
          integers, and is 1d or a column vector.
        * 'continuous-multioutput': `y` is a 2d array of floats that are
          not all integers, and both dimensions are of size > 1.
        * 'binary': `y` contains <= 2 discrete values and is 1d or a column
          vector.
        * 'multiclass': `y` contains more than two discrete values, is not a
          sequence of sequences, and is 1d or a column vector.
        * 'multiclass-multioutput': `y` is a 2d array that contains more
          than two discrete values, is not a sequence of sequences, and both
          dimensions are of size > 1.
        * 'multilabel-indicator': `y` is a label indicator matrix, an array
          of two dimensions with at least two columns, and at most 2 unique
          values.
        * 'unknown': `y` is array-like but none of the above, such as a 3d
          array, sequence of sequences, or an array of non-sequence objects.

    Examples
    --------
    >>> import numpy as np
    >>> type_of_target([0.1, 0.6])
    'continuous'
    >>> type_of_target([1, -1, -1, 1])
    'binary'
    >>> type_of_target(['a', 'b', 'a'])
    'binary'
    >>> type_of_target([1.0, 2.0])
    'binary'
    >>> type_of_target([1, 0, 2])
    'multiclass'
    >>> type_of_target([1.0, 0.0, 3.0])
    'multiclass'
    >>> type_of_target(['a', 'b', 'c'])
    'multiclass'
    >>> type_of_target(np.array([[1, 2], [3, 1]]))
    'multiclass-multioutput'
    >>> type_of_target([[1, 2]])
    'multiclass-multioutput'
    >>> type_of_target(np.array([[1.5, 2.0], [3.0, 1.6]]))
    'continuous-multioutput'
    >>> type_of_target(np.array([[0, 1], [1, 1]]))
    'multilabel-indicator'
    """
    valid = ((isinstance(y, (Sequence, spmatrix)) or hasattr(y, '__array__'))
             and not isinstance(y, string_types))

    if not valid:
        raise ValueError('Expected array-like (array or non-string sequence), '
                         'got %r' % y)

    if is_multilabel(y):
        return 'multilabel-indicator'

    try:
        y = np.asarray(y)
    except ValueError:
        # Known to fail in numpy 1.3 for array of arrays
        return 'unknown'

    # The old sequence of sequences format
    try:
        if (not hasattr(y[0], '__array__') and isinstance(y[0], Sequence)
                and not isinstance(y[0], string_types)):
            raise ValueError('You appear to be using a legacy multi-label data'
                             ' representation. Sequence of sequences are no'
                             ' longer supported; use a binary array or sparse'
                             ' matrix instead.')
    except IndexError:
        pass

    # Invalid inputs
    if y.ndim > 2 or (y.dtype == object and len(y) and
                      not isinstance(y.flat[0], string_types)):
        return 'unknown'  # [[[1, 2]]] or [obj_1] and not ["label_1"]

    if y.ndim == 2 and y.shape[1] == 0:
        return 'unknown'  # [[]]

    if y.ndim == 2 and y.shape[1] > 1:
        suffix = "-multioutput"  # [[1, 2], [1, 2]]
    else:
        suffix = ""  # [1, 2, 3] or [[1], [2], [3]]

    # check float and contains non-integer float values
    if y.dtype.kind == 'f' and np.any(y != y.astype(int)):
        # [.1, .2, 3] or [[.1, .2, 3]] or [[1., .2]] and not [1., 2., 3.]
        return 'continuous' + suffix

    if (len(np.unique(y)) > 2) or (y.ndim >= 2 and len(y[0]) > 1):
        return 'multiclass' + suffix  # [1, 2, 3] or [[1., 2., 3]] or [[1, 2]]
    else:
        return 'binary'  # [1, 2] or [["a"], ["b"]]

def _check_partial_fit_first_call(clf, classes=None):
    """Private helper function for factorizing common classes param logic

    Estimators that implement the ``partial_fit`` API need to be provided with
    the list of possible classes at the first call to partial_fit.

    Subsequent calls to partial_fit should check that ``classes`` is still
    consistent with a previous value of ``clf.classes_`` when provided.

    This function returns True if it detects that this was the first call to
    ``partial_fit`` on ``clf``. In that case the ``classes_`` attribute is also
    set on ``clf``.

    """
    if getattr(clf, 'classes_', None) is None and classes is None:
        raise ValueError("classes must be passed on the first call "
                         "to partial_fit.")

    elif classes is not None:
        if getattr(clf, 'classes_', None) is not None:
            if not array_equal(clf.classes_, unique_labels(classes)):
                raise ValueError(
                    "`classes=%r` is not the same as on last call "
                    "to partial_fit, was: %r" % (classes, clf.classes_))

        else:
            # This is the first call to partial_fit
            clf.classes_ = unique_labels(classes)
            return True

    # classes is None and clf.classes_ has already previously been set:
    # nothing to do
    return False

def class_distribution(y, sample_weight=None):
    """Compute class priors from multioutput-multiclass target data

    Parameters
    ----------
    y : array like or sparse matrix of size (n_samples, n_outputs)
        The labels for each example.

    sample_weight : array-like of shape = (n_samples,), optional
        Sample weights.

    Returns
    -------
    classes : list of size n_outputs of arrays of size (n_classes,)
        List of classes for each column.

    n_classes : list of integers of size n_outputs
        Number of classes in each column

    class_prior : list of size n_outputs of arrays of size (n_classes,)
        Class distribution of each column.

    """
    classes = []
    n_classes = []
    class_prior = []

    n_samples, n_outputs = y.shape

    if issparse(y):
        y = y.tocsc()
        y_nnz = np.diff(y.indptr)

        for k in range(n_outputs):
            col_nonzero = y.indices[y.indptr[k]:y.indptr[k + 1]]
            # separate sample weights for zero and non-zero elements
            if sample_weight is not None:
                nz_samp_weight = np.asarray(sample_weight)[col_nonzero]
                zeros_samp_weight_sum = (np.sum(sample_weight) -
                                         np.sum(nz_samp_weight))
            else:
                nz_samp_weight = None
                zeros_samp_weight_sum = y.shape[0] - y_nnz[k]

            classes_k, y_k = np.unique(y.data[y.indptr[k]:y.indptr[k + 1]],
                                       return_inverse=True)
            class_prior_k = bincount(y_k, weights=nz_samp_weight)

            # An explicit zero was found, combine its weight with the weight
            # of the implicit zeros
            if 0 in classes_k:
                class_prior_k[classes_k == 0] += zeros_samp_weight_sum

            # If an there is an implicit zero and it is not in classes and
            # class_prior, make an entry for it
            if 0 not in classes_k and y_nnz[k] < y.shape[0]:
                classes_k = np.insert(classes_k, 0, 0)
                class_prior_k = np.insert(class_prior_k, 0,
                                          zeros_samp_weight_sum)

            classes.append(classes_k)
            n_classes.append(classes_k.shape[0])
            class_prior.append(class_prior_k / class_prior_k.sum())
    else:
        for k in range(n_outputs):
            classes_k, y_k = np.unique(y[:, k], return_inverse=True)
            classes.append(classes_k)
            n_classes.append(classes_k.shape[0])
            class_prior_k = bincount(y_k, weights=sample_weight)
            class_prior.append(class_prior_k / class_prior_k.sum())

    return (classes, n_classes, class_prior)


import requests
import pandas as pd
import os.path
import resources.constants
import json
from pandas.io.json import json_normalize
from utils.data_utils import get_ordinal_number

"""
This script aggregates data from multiple endpoints and returns a single .json file containing all data
used in the displacement tracker project.

Scheduling this script would mean that the /displacement_tracker endpoint always returned the latest data
contained within the Humanitarian Data Service API.
"""

# For development
#ROOT = 'http://localhost:5000'

# For live
ROOT = 'http://ec2-34-200-18-111.compute-1.amazonaws.com'

# Set year for country-level funding data
FUNDING_YEAR = 2016

# Define all endpoints
URL_POPULATIONS_REFUGEELIKE_ASYLUM = '/populations/refugeelike/asylum/index'
URL_POPULATIONS_REFUGEELIKE_ORIGIN = '/populations/refugeelike/origin/index'
URL_INDICATORS_GNI = '/indicators/gni/index'
URL_PLANS_PROGRESS = '/funding/plans/progress/index'
URL_POPULATION = '/populations/totals/index'
URL_FRAGILE_STATE = '/fragility/fragile-state-index/index'
URL_NEEDS = '/needs/plans/index'
URL_FUNDING_DEST_COUNTRY = '/funding/countries/destination/index/{}'.format(FUNDING_YEAR)
URL_FUNDING_DEST_DONORS = '/funding/countries/donors/index'

# Define path for raw country names data
country_names_path = os.path.join(resources.constants.EXAMPLE_RAW_DATA_PATH, 'UNSD Methodology.csv')

# Define path for relatable geography populations data
relatable_population_path = os.path.join(resources.constants.EXAMPLE_DERIVED_DATA_PATH, '2017_relatable_population_rankings.csv')

# Define path for stories of displacement
displacement_stories_path = os.path.join(resources.constants.EXAMPLE_DERIVED_DATA_PATH, 'stories_of_displacement_links.csv')

# Create a blank dictionary to store metadata for each field
metadata_dict = {}

def merge_data(
        funding_year = FUNDING_YEAR,
        country_names_path=country_names_path,
        relatable_population_path=relatable_population_path,
        displacement_stories_path=displacement_stories_path,
        url_populations_refugeelike_asylum=(ROOT + URL_POPULATIONS_REFUGEELIKE_ASYLUM),
        url_populations_refugeelike_origin=(ROOT + URL_POPULATIONS_REFUGEELIKE_ORIGIN),
        url_indicators_gni=(ROOT + URL_INDICATORS_GNI),
        url_plans_progress=(ROOT + URL_PLANS_PROGRESS),
        url_population=(ROOT + URL_POPULATION),
        url_fragile_state=(ROOT + URL_FRAGILE_STATE),
        url_needs=(ROOT + URL_NEEDS),
        url_funding_dest_country=(ROOT + URL_FUNDING_DEST_COUNTRY),
        url_funding_dest_donors=(ROOT + URL_FUNDING_DEST_DONORS)
    ):

    ####################  COUNTRY NAMES ####################
    # Get the data from .csv
    df_country_names = pd.read_csv(country_names_path, encoding='utf-8')

    # Select relevant fields
    df_country_names = df_country_names[[
        'Country or Area',
        'ISO-alpha3 Code'
    ]]

    # Add Taiwan
    df_country_names.loc[-1] = ["Taiwan", "TWN"]

    # Drop null values
    df_country_names = df_country_names.dropna()

    # Set country code to be the index
    df_country_names = df_country_names.set_index('ISO-alpha3 Code')

    # Rename fields
    df_country_names.rename(columns={'Country or Area': 'Country'}, inplace=True)

    ####################  DISPLACEMENT STORIES ####################
    # Get the data from .csv
    df_displacement_stories = pd.read_csv(displacement_stories_path, encoding='utf-8')

    # Set country code to be the index
    df_displacement_stories = df_displacement_stories.set_index('countryCode')

    # Select relevant fields
    df_displacement_stories = df_displacement_stories[[
        'storyTitle', 'storySource',
        'storyTagLine', 'storyURL'
    ]]

    # Drop null values
    df_displacement_stories = df_displacement_stories.dropna()

    # Add metadata for each field to overall metadata dictionary
    for column in df_displacement_stories.columns:
        metadata_dict[column] = {}

    ####################  POPULATIONS ####################
    # Get the data from the API
    population_data = requests.get(url_population).json()

    # Extract metadata
    if 'metadata' in population_data:
        population_metadata = population_data['metadata']
    else:
        population_metadata = {}

    # Build dataframe
    df_population = pd.DataFrame(population_data['data']).T

    # Select relevant fields
    df_population = df_population[[
        'PopTotal'
    ]]

    # Rename fields
    df_population.rename(columns={'PopTotal': 'Population'}, inplace=True)

    # Drop null values
    df_population = df_population.dropna()

    # Add metadata for each field to overall metadata dictionary
    for column in df_population.columns:
        metadata_dict[column] = population_metadata

    ####################  FRAGILE STATE ####################
    # Get the data from the API
    fragile_state_data = requests.get(url_fragile_state).json()

    # Extract metadata
    if 'metadata' in fragile_state_data:
        fragile_state_metadata = fragile_state_data['metadata']
    else:
        fragile_state_metadata = {}

    # Build a dataframe
    df_fragile_state = pd.DataFrame(fragile_state_data['data']).T

    # Select relevant fields
    df_fragile_state = df_fragile_state[[
        'Total', 'Rank'
    ]]

    # Rename fields
    df_fragile_state.rename(columns={'Total': 'Fragile State Index Score',
                                     'Rank': 'Fragile State Index Rank'}, inplace=True)

    # Drop null values
    df_fragile_state = df_fragile_state.dropna()

    # Add metadata for each field to overall metadata dictionary
    for column in df_fragile_state.columns:
        metadata_dict[column] = fragile_state_metadata

    ####################  POPULATIONS_REFUGEELIKE_ASYLUM ####################
    # Get the data from the API
    populations_refugeelike_asylum_data = requests.get(url_populations_refugeelike_asylum).json()

    # Extract metadata
    if 'metadata' in populations_refugeelike_asylum_data:
        populations_refugeelike_asylum_metadata = populations_refugeelike_asylum_data['metadata']
    else:
        populations_refugeelike_asylum_metadata = {}

    # Build a dataframe
    df_populations_refugeelike_asylum = pd.DataFrame(populations_refugeelike_asylum_data['data']).T

    # Select relevant fields
    df_populations_refugeelike_asylum = df_populations_refugeelike_asylum[[
        'Total population of concern', 'Total Refugee and people in refugee-like situations',
        'IDPs protected/assisted by UNHCR, incl. people in IDP-like situations','Asylum-seekers'
    ]]

    # Rename fields
    df_populations_refugeelike_asylum.rename(columns={
        'IDPs protected/assisted by UNHCR, incl. people in IDP-like situations': 'IDPs protected/assisted by UNHCR',
        'Asylum-seekers': 'Asylum-seekers (asylum)'
    }, inplace=True)

    # Add field to rank total total population of concern
    df_populations_refugeelike_asylum['Rank of total population of concern'] = df_populations_refugeelike_asylum[
        'Total population of concern'].rank(ascending=False, method='min').astype(int)

    # Add field to add refugees and asylum-seekers
    df_populations_refugeelike_asylum['Total refugees and asylum-seekers (asylum)'] = df_populations_refugeelike_asylum[
        'Total Refugee and people in refugee-like situations'] + df_populations_refugeelike_asylum['Asylum-seekers (asylum)']

    # Drop null values
    df_populations_refugeelike_asylum = df_populations_refugeelike_asylum.dropna()

    # Add metadata for each field to overall metadata dictionary
    for column in df_populations_refugeelike_asylum.columns:
        metadata_dict[column] = populations_refugeelike_asylum_metadata

    ####################  POPULATIONS_REFUGEELIKE_ORIGIN ####################
    # Get the data from the API
    populations_refugeelike_origin_data = requests.get(url_populations_refugeelike_origin).json()

    # Extract metadata
    if 'metadata' in populations_refugeelike_origin_data:
        populations_refugeelike_origin_metadata = populations_refugeelike_origin_data['metadata']
    else:
        populations_refugeelike_origin_metadata = {}

    # Build a dataframe
    df_populations_refugeelike_origin = pd.DataFrame(populations_refugeelike_origin_data['data']).T

    # Select relevant fields
    df_populations_refugeelike_origin = df_populations_refugeelike_origin[[
        'Total Refugee and people in refugee-like situations', 'Asylum-seekers'
    ]]

    # Rename fields
    df_populations_refugeelike_origin.rename(columns={
        'Total Refugee and people in refugee-like situations': 'Total refugees who have fled from country',
        'Asylum-seekers': 'Asylum-seekers (origin)'
    }, inplace=True)

    # Add field to add refugees and asylum-seekers
    df_populations_refugeelike_origin['Total refugees and asylum-seekers (origin)'] = df_populations_refugeelike_origin[
        'Total refugees who have fled from country'] + df_populations_refugeelike_origin['Asylum-seekers (origin)']

    # Drop null values
    df_populations_refugeelike_origin = df_populations_refugeelike_origin.dropna()

    # Add metadata for each field to overall metadata dictionary
    for column in df_populations_refugeelike_origin.columns:
        metadata_dict[column] = populations_refugeelike_origin_metadata

    ####################  INDICATORS GNI ####################
    # Get the data from the API
    indicators_gni_data = requests.get(url_indicators_gni).json()

    # Extract metadata
    if 'metadata' in indicators_gni_data:
        indicators_gni_metadata = indicators_gni_data['metadata']
    else:
        indicators_gni_metadata = {}

    # Build a dataframe
    df_indicators_gni = pd.DataFrame(indicators_gni_data['data']).T

    # Select relevant fields
    df_indicators_gni = df_indicators_gni[[
        '2015'
    ]]

    # Rename fields
    df_indicators_gni.rename(columns={'2015': 'GDP Per Capita'}, inplace=True)

    # Drop null values
    df_indicators_gni = df_indicators_gni.dropna()

    # Add metadata for each field to overall metadata dictionary
    for column in df_indicators_gni.columns:
        metadata_dict[column] = indicators_gni_metadata

    ####################  PLANS PROGRESS ####################
    # Get the data from the API
    plans_progress_data = requests.get(url_plans_progress).json()

    # Extract metadata
    if 'metadata' in plans_progress_data:
        plans_progress_metadata = plans_progress_data['metadata']
    else:
        plans_progress_metadata = {}

    # Build a dataframe
    df_plans_progress = pd.DataFrame(plans_progress_data['data']).T

    # Select relevant fields
    df_plans_progress = df_plans_progress[[
        'appealFunded', 'revisedRequirements', 'neededFunding'
    ]]

    # Rename fields
    df_plans_progress.rename(columns={'appealFunded': 'Appeal funds committed to date',
                                      'revisedRequirements': 'Appeal funds requested',
                                      'neededFunding': 'Appeal funds still needed'}, inplace=True)

    df_plans_progress['Appeal percent funded'] = df_plans_progress['Appeal funds committed to date']/df_plans_progress['Appeal funds requested']

    # Drop null values
    df_plans_progress = df_plans_progress.dropna()

    # Add metadata for each field to overall metadata dictionary
    for column in df_plans_progress.columns:
        metadata_dict[column] = plans_progress_metadata

    # Add an FTS data as-of date so it can be included in the .csv data dump
    df_plans_progress['FTS funding data as-of date'] = plans_progress_data['metadata']['source_data']

    ######## FUNDING BY DESTINATION COUNTRY ############
    #Get the data from the API
    funding_dest_country_data = requests.get(url_funding_dest_country).json()

    # Extract metadata
    if 'metadata' in funding_dest_country_data:
        funding_dest_country_metadata = funding_dest_country_data['metadata']
    else:
        funding_dest_country_metadata = {}

    # Build a dataframe
    df_funding_dest_country = pd.DataFrame(funding_dest_country_data['data']).T

    # Select relevant fields
    df_funding_dest_country = df_funding_dest_country[[
        'totalFunding'
    ]]

    # Keep only records where totalFunding > 0
    df_funding_dest_country = df_funding_dest_country[df_funding_dest_country['totalFunding'] > 0]

    # Rename fields
    df_funding_dest_country.rename(columns={'totalFunding': 'Humanitarian aid received'},
                                   inplace=True)

    # Add field to rank total total population of concern
    df_funding_dest_country['Rank of humanitarian aid received'] = df_funding_dest_country[
        'Humanitarian aid received'].rank(ascending=False, method='min').astype(int)

    # Drop null values
    df_funding_dest_country = df_funding_dest_country.dropna()

    # Add metadata for each field to overall metadata dictionary
    for column in df_funding_dest_country.columns:
        metadata_dict[column] = funding_dest_country_metadata

    ################## TOP 5 DONORS TO EACH DESTINATION COUNTRY ###################
    #Get the data from the API
    funding_dest_donors_data = requests.get(url_funding_dest_donors).json()

    # Extract metadata
    if 'metadata' in funding_dest_donors_data:
        funding_dest_donors_metadata = funding_dest_donors_data['metadata']
    else:
        funding_dest_donors_metadata = {}

    # Build a dataframe
    df_funding_dest_donors = json_normalize(funding_dest_donors_data['data']).T
    #df_funding_dest_donors = pd.DataFrame(funding_dest_donors_data['data']).T

    df_funding_dest_donors.columns = (['Top 5 Donors'])

    # Add metadata for each field to overall metadata dictionary
    for column in df_funding_dest_donors.columns:
        metadata_dict[column] = funding_dest_donors_metadata

    ####################  NEEDS ####################
    # Get the data from the API
    needs_data = requests.get(url_needs).json()

    # Extract metadata
    if 'metadata' in needs_data:
        needs_metadata = needs_data['metadata']
    else:
        needs_metadata = {}

    # Build a dataframe
    df_needs = pd.DataFrame(needs_data['data']).T

    # Exclude rows where country code is missing
    df_needs = df_needs.drop('null')

    # Select relevant fields
    df_needs = df_needs[[
        'inNeedTotal', 'inNeedHealth', 'inNeedEducation',
        'inNeedFoodSecurity', 'inNeedProtection', 'sourceURL',
        'inNeedShelter-CCCM-NFI', 'inNeedWASH', 'sourceType'
    ]]

    # Rename fields
    df_needs.rename(columns={'inNeedTotal': 'Total people in need',
                             'inNeedHealth': 'People in need of health support',
                             'inNeedEducation': 'Children in need of education',
                             'inNeedFoodSecurity': 'People who are food insecure',
                             'inNeedProtection': 'People in need of protection',
                             'inNeedShelter-CCCM-NFI': 'People in need of shelter',
                             'inNeedWASH': 'People in need of water, sanitization & hygiene',
                             'sourceURL': 'Source of needs data',
                             'sourceType': 'Source type of needs data'
                             }, inplace=True)

    # Add metadata for each field to overall metadata dictionary
    for column in df_needs.columns:
        metadata_dict[column] = needs_metadata

    ######## FIND PLACES WITH SIMILAR POPULATIONS TO PEOPLE IN NEED ########

    # Get the relateable populations data from .csv
    df_relatable_populations = pd.read_csv(relatable_population_path)
    df_relatable_populations['Population'] = df_relatable_populations[[
        'Population - World Bank (2015)','Population - UNFPA (2016)'
    ]].max(axis=1)

    df_relatable_populations = df_relatable_populations[['City, State, Country','Population']].dropna()

    def find_nearest_place_population(reference_value):

        if reference_value:
            nearest_row = df_relatable_populations.iloc[(df_relatable_populations['Population']- reference_value).abs().argsort()[0]]
            nearest_population = nearest_row['Population']
        else:
            nearest_population = 0.00

        return nearest_population

    def find_nearest_place(reference_value):

        if reference_value:
            nearest_row = df_relatable_populations.iloc[(df_relatable_populations['Population']- reference_value).abs().argsort()[0]]
            nearest_place = nearest_row['City, State, Country']
        else:
            nearest_place = ''

        return nearest_place

    df_needs['Place with similar population as people in need'] = df_needs['Total people in need'].apply(
        find_nearest_place)
    # Add metadata
    metadata_dict['Place with similar population as people in need'] = {}

    df_needs['Population of place with similar population'] = df_needs['Total people in need'].apply(
        find_nearest_place_population)
    # Add metadata
    metadata_dict['Population of place with similar population'] = {}

    ####################  SAMPLE CLUSTERS ####################

    # Build a dataframe
    # df_clusters = pd.read_json('sample_clusters.json').T
    # df_clusters = df_clusters[['clusters']]

    ################# COMBINE ALL DATA ##############

    # Make a list of all dataframes
    all_dataframes = [
        df_country_names,
        df_populations_refugeelike_asylum,
        df_indicators_gni,
        df_plans_progress,
        df_population,
        df_fragile_state,
        df_needs,
        df_funding_dest_country,
        df_funding_dest_donors,
        df_displacement_stories,
        df_populations_refugeelike_origin
        #   df_clusters
    ]

    df_final = pd.concat(all_dataframes, axis=1)

    # Add calculation for displaced people as a ratio of total population
    df_final['Population of concern per 1000 population'] = (df_final['Total population of concern'] / df_final[
        'Population'])*1000
    # And metadata
    metadata_dict['Population of concern per 1000 population'] = {}
    metadata_dict['Population of concern per 1000 population']['Calculation'] = '(Total population of concern / Population) * 1000'

    # Add calculation for displaced people per million GDP
    df_final['Population of concern per million GDP'] = ((df_final['Total population of concern'] * 1000000) / (df_final[
        'GDP Per Capita'] * df_final['Population']))
    # And metadata
    metadata_dict['Population of concern per million GDP'] = {}
    metadata_dict['Population of concern per million GDP']['Calculation'] = '(Total population of concern] * 1000000) / (GDP Per Capita * Population)'

    # Add field to specify whether country has current humanitarian appeal in FTS
    df_final['Country has current appeal'] = df_final['Appeal funds requested'].notnull()
    # And metadata
    metadata_dict['Country has current appeal'] = {}
    metadata_dict['Country has current appeal']['Calculation'] = 'Is Appeal funds requested not null'

    # Make the ranked variables ordinal

    def get_ordinal_number(value):
        try:
            value = int(value)
        except ValueError:
            return value

        if value % 100 // 10 != 1:
            if value % 10 == 1:
                ordval = u"%d%s" % (value, "st")
            elif value % 10 == 2:
                ordval = u"%d%s" % (value, "nd")
            elif value % 10 == 3:
                ordval = u"%d%s" % (value, "rd")
            else:
                ordval = u"%d%s" % (value, "th")
        else:
            ordval = u"%d%s" % (value, "th")

        return ordval

    df_final['Rank of total population of concern'] = df_final['Rank of total population of concern'].apply(
        get_ordinal_number)

    df_final['Rank of humanitarian aid received'] = df_final['Rank of humanitarian aid received'].apply(
        get_ordinal_number)

    ################## STRUCTURE DICTIONARY ##################

    # Clean up NaN values
    df_final = df_final.fillna('')

    # Transform dataframe to dictionary
    df_as_dict = df_final.to_dict(orient='index')

    # Define field names for each strand
    strand_01_fields = ['Appeal funds still needed', 'Appeal funds requested', 'Appeal funds committed to date',
                        'Appeal percent funded', 'Source of needs data', 'Source type of needs data',
                        'Total people in need', 'Place with similar population as people in need',
                        'Population of place with similar population']
    strand_02_fields = ['Population of concern per 1000 population', 'Fragile State Index Score',
                        'Total population of concern',
                        'IDPs protected/assisted by UNHCR',
                        'GDP Per Capita',
                        'Total refugees and asylum-seekers (asylum)',
                        'Total refugees and asylum-seekers (origin)']
    strand_03_fields = ['Humanitarian aid received', 'Appeal funds requested', 'Appeal percent funded',
                        'Rank of total population of concern', 'Rank of humanitarian aid received']

    needs_fields = ['People in need of health support','Children in need of education',
                    'People who are food insecure','People in need of protection','People in need of shelter',
                    'People in need of water, sanitization & hygiene']

    story_fields = ['storyTitle', 'storySource', 'storyTagLine', 'storyURL']

    # For every object, get / group the values by strand
    data = {}
    for x in df_as_dict.keys():

        # Create an empty dict
        country_dict = {}

        # Populate the dict with those value that don't require nesting
        country_dict['Country'] = df_as_dict[x]['Country']
        country_dict['Fragile State Index Rank'] = df_as_dict[x]['Fragile State Index Rank']
        country_dict['Country has current appeal'] = df_as_dict[x]['Country has current appeal']

        # Populate the dict with story fields
        story_fields_dict = {}
        if df_as_dict[x]['storyURL']:
            for field in story_fields:
                story_fields_dict[field] = (df_as_dict[x][field])
        country_dict['Displacement_story'] = story_fields_dict

        # Populate the dict with strand 1 data if the country has a current appeal
        strand_01_dict = {}
        if df_as_dict[x]['Country has current appeal']:
            strand_01_dict['Needs_Data'] = {}
            for names_01 in strand_01_fields:
                strand_01_dict[names_01] = (df_as_dict[x][names_01])
            for name in needs_fields:
                if df_as_dict[x][name] != '':
                    strand_01_dict['Needs_Data'][name] = (df_as_dict[x][name])
        country_dict['Strand_01_Needs'] = strand_01_dict

        # Populate the dict with strand 2 data
        strand_02_dict = {}
        for names_02 in strand_02_fields:
            strand_02_dict[names_02] = (df_as_dict[x][names_02])
        country_dict['Strand_02_People'] = strand_02_dict

        # Populate the dict with strand 3 data
        strand_03_dict = {}
        strand_03_dict['Top 5 donors of humanitarian aid'] = []
        for names_03 in strand_03_fields:
            strand_03_dict[names_03] = (df_as_dict[x][names_03])
        if df_as_dict[x]['Top 5 Donors']:
            strand_03_dict['Top 5 donors of humanitarian aid'] = df_as_dict[x]['Top 5 Donors']
        country_dict['Strand_03_Aid'] = strand_03_dict

        # Add the country dict to the data dict
        data[x] = country_dict

    # Add World totals
    # Create an empty dict
    world_dict = {}

    # Populate the dict with aggregated strand 1 data
    strand_01_dict = {}
    strand_01_dict['Needs_Data'] = {}
    strand_01_dict['Total people in need'] = df_needs['Total people in need'].sum()
    strand_01_dict['Count of current crises with people in need'] = df_needs['Total people in need'].count()
    strand_01_dict['Place with similar population as people in need'] = find_nearest_place(
        df_needs['Total people in need'].sum()
    )
    strand_01_dict['Population of place with similar population'] = find_nearest_place_population(
        df_needs['Total people in need'].sum()
    )
    for name in needs_fields:
        strand_01_dict['Needs_Data'][name] = df_needs[name].sum()
    world_dict['Strand_01_Needs'] = strand_01_dict

    # Add the world dict to the data dict
    data['WORLD'] = world_dict

    # Create the metadata dict
    metadata = {}

    # Populate the dict with those value that don't require nesting
    #metadata['Country'] = metadata_dict['Country']
    metadata['Fragile State Index Rank'] = metadata_dict['Fragile State Index Rank']
    metadata['Country has current appeal'] = metadata_dict['Country has current appeal']

    # Populate the dict with story fields
    story_fields_dict = {}
    if metadata_dict['storyURL']:
        for field in story_fields:
            story_fields_dict[field] = (metadata_dict[field])
    metadata['Displacement_story'] = story_fields_dict

    # Populate the dict with strand 1 data if the country has a current appeal
    strand_01_dict = {}
    strand_01_dict['Needs_Data'] = {}
    for names_01 in strand_01_fields:
        strand_01_dict[names_01] = (metadata_dict[names_01])
    metadata['Strand_01_Needs'] = strand_01_dict

    # Populate the dict with strand 2 data
    strand_02_dict = {}
    for names_02 in strand_02_fields:
        strand_02_dict[names_02] = (metadata_dict[names_02])
    metadata['Strand_02_People'] = strand_02_dict

    # Populate the dict with strand 3 data
    strand_03_dict = {}
    strand_03_dict['Top 5 donors of humanitarian aid'] = []
    for names_03 in strand_03_fields:
        strand_03_dict[names_03] = (metadata_dict[names_03])
    if metadata_dict['Top 5 Donors']:
        strand_03_dict['Top 5 donors of humanitarian aid'] = metadata_dict['Top 5 Donors']
    metadata['Strand_03_Aid'] = strand_03_dict

    # At the higher level, structure the json with 'data' and 'metadata'
    final_json = {
        'data': data,
        'metadata': metadata
    }

    return final_json, metadata, df_final

def run():
    print 'Pulling and merging data'
    final_json, metadata, final_csv = merge_data()

    print 'Writing Combined JSON file'
    with open(os.path.join(resources.constants.EXAMPLE_DERIVED_DATA_PATH, 'displacement_tracker.json'), 'w') as outfile:
        json.dump(final_json, outfile, indent=4, separators=(',', ': '), ensure_ascii=True, sort_keys=True)

    print 'Writing Combined JSON metadata file'
    with open(os.path.join(resources.constants.EXAMPLE_DERIVED_DATA_PATH, 'displacement_tracker_metadata.json'), 'w') as outfile:
        json.dump(metadata, outfile, indent=4, separators=(',', ': '), ensure_ascii=True, sort_keys=True)

    print 'Writing Combined CSV file'
    final_csv.to_csv(os.path.join(resources.constants.EXAMPLE_DERIVED_DATA_PATH, 'displacement_tracker.csv'), index_label='CountryCode', encoding='utf-8')

if __name__ == "__main__":
    run()