"""
Unit tests for join.py
----------------------

- Tests are based on splitting a field along the 'layers axis' and joining it back together.
- Fields passed to unite must fulfill a set of basic requirements, which are also tested.
- Optional checks on the fields are not tested here because they have their own dedicated tests.
"""
import random
import pytest
from canopy.tests.test_data.registry import load_test_data
from canopy.util.join import join
from canopy.core.field import Field
from canopy.core.grid.registry import get_grid

TEST_DATASET_LONLAT = 'anpp_spain_1990_2010.out.gz'
TEST_DATASET_SITES = 'aaet_global_sites.out.gz'

def test_disjoint_layers_same_indices():
    """Join fields with equal indices and disjoint layers"""

    field = load_test_data(TEST_DATASET_LONLAT)

    layers = field.layers
    i1, i2 = sorted(random.sample(range(len(layers)), 2))
    l1, l2, l3 = layers[:i1], layers[i1:i2], layers[i2:]
    f1, f2, f3 = field[l1], field[l2], field[l3]

    field_joined = join([f1, f2, f3])

    assert field.data.equals(field_joined.data)

    
def test_disjoint_layers_and_indices():
    """
    Join fields with disjoint layers and indices

    Check for number of NaNs in joined field
    """

    field = load_test_data(TEST_DATASET_LONLAT)
    dlon = field.grid.dlon
    dlat = field.grid.dlat

    # Create fields with disjoint indices
    i1, i2 = sorted(random.sample(range(len(field.data)), 2))
    df1, df2, df3 = field.data.iloc[:i1], field.data.iloc[i1:i2], field.data.iloc[i2:]
    r1, r2, r3 = len(df1), len(df2), len(df3)
    f1, f2, f3 = ( Field(df, get_grid(field.grid.grid_type).from_frame(df, dlon=dlon, dlat=dlat)) for df in [df1, df2, df3] )

    # Select disjoint layers
    layers = field.layers
    i1, i2 = sorted(random.sample(range(len(layers)), 2))
    l1, l2, l3 = layers[:i1], layers[i1:i2], layers[i2:]
    f1, f2, f3 = f1[l1], f2[l2], f3[l3]

    number_of_nans = len(l1)*(r2 + r3) + len(l2)*(r1 + r3) + len(l3)*(r1 + r2)

    field_joined = join([f1, f2, f3])
    
    assert number_of_nans == field_joined.data.isna().sum().sum()


def test_different_frequencies():
    """Returned field has only overlapping layers from the pieces."""

    field_1Y = load_test_data(TEST_DATASET_LONLAT)
    field_5Y = field_1Y.reduce_time('av', '5Y')

    with pytest.raises(ValueError, match="All fields must have the same time frequency.") as v_error:
        join([field_1Y, field_5Y])


def test_different_grid_types():

    field_lonlat = load_test_data(TEST_DATASET_LONLAT)
    field_sites = load_test_data(TEST_DATASET_SITES)

    with pytest.raises(ValueError, match="Field grids must all be of the same type.") as v_error:
        join([field_sites, field_lonlat])


def test_incompatible_grid_types():

    field1 = load_test_data(TEST_DATASET_LONLAT)

    df = field1.data.reset_index()
    df['lon'] += 0.1
    df = df.set_index(['lon', 'lat', 'time'])
    grid = get_grid(field1.grid.grid_type).from_frame(df)
    field2 = Field(df, grid)

    df = field2.data.reset_index()
    df['lat'] += 0.1
    df = df.set_index(['lon', 'lat', 'time'])
    grid = get_grid(field1.grid.grid_type).from_frame(df)
    field3 = Field(df, grid)

    with pytest.raises(ValueError, match="Field grids must all be mutually compatible.") as v_error:
        join([field1, field1, field2, field3])


def test_overlapping_layers():

    field = load_test_data(TEST_DATASET_LONLAT)

    layers = field.layers
    i1, i2 = sorted(random.sample(range(len(layers)), 2))
    l1, l2, l3 = layers[:i1], layers[i1-i1//2:i2], layers[i2-i2//2:]
    f1, f2, f3 = field[l1], field[l2], field[l3]

    with pytest.raises(ValueError, match="Fields have common layers.") as v_error:
        join([f1, f2, f3])
