#!/usr/bin/env python

import pytest
import pandas as pd
import numpy as np
from mapies.grids.monarch import RegularGrid, RotatedGrid, IrregularRotatedGrid


@pytest.mark.parametrize(
    "test_input,expected", 
    [
        ((np.array([-10, -10, -10]), np.array([-10, -10, -10]), np.array([1, 2, 3])),
        (np.array([-10., 0., 0., 0.]), np.array([-10., 0., 0., 0.]), np.array([2., 0., 0., 0.]), np.array([[3., 0., 0., 0.]]))),
    ]
)
def test_regular_grid_aggregation(test_input, expected):
    r = RegularGrid(dlon=180, dlat=90, west=-180, south=-90)
    agg = r.aggregate(*test_input)
    assert agg[0].all() == expected[0].all()
    assert agg[1].all() == expected[1].all()
    assert agg[2].all() == expected[2].all()
    assert agg[3].all() == expected[3].all()


@pytest.mark.parametrize(
    "test_input,expected", 
    [
        ((np.array([0, 0, 0]), np.array([0, 0, 0]), np.array([1, 2, 3])), 
        (np.array([0.]), np.array([0.]), np.array([0.]), np.array([0.]), np.array([2.]), np.array([3]))),
    ]
)
def test_rotated_grid_aggregation(test_input, expected):
    r = RotatedGrid(dlon=1, dlat=1, west=-1, south=-1, centre_lon=0, centre_lat=0)
    agg = r.aggregate(*test_input)
    assert agg[0].all() == expected[0].all()
    assert agg[1].all() == expected[1].all()
    assert agg[2].all() == expected[2].all()
    assert agg[3].all() == expected[3].all()



@pytest.mark.parametrize(
    "test_input,expected", 
    [
        ((np.array([0, 0, 0]), np.array([0, 0, 0]), np.array([1, 2, 3])), 
        (np.array([0.]), np.array([0.]), np.array([0.]), np.array([0.]), np.array([2.]), np.array([3]))),
    ]
)
def test_irregular_rotated_grid_aggregation(test_input, expected):
    r = IrregularRotatedGrid(dlon=1, dlat=1, west=-1, south=-1, centre_lon=0, centre_lat=0)
    agg = r.aggregate(*test_input)
    assert agg[0].all() == expected[0].all()
    assert agg[1].all() == expected[1].all()
    assert agg[2].all() == expected[2].all()
    assert agg[3].all() == expected[3].all()