#!/usr/bin/env python

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


@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()

@pytest.mark.parametrize(
    "test_input,expected",
    [
        ((0, 0, 20, 35), (-23.95680324, -32.61460715)),
        ((10, 10, 20, 35), (-10.82847335, -24.45866947)),
        ((-10, -10, 20, 35), (-39.42034458, -39.15567129)),
        ((-100, -50, 20, 35), (-141.61239392, -26.30586515))
    ]
)
def test_geo_to_rot(test_input, expected):
    assert (round(geo_to_rot(*test_input)[0], 8), round(geo_to_rot(*test_input)[1], 8))  == expected