"""
Tests for vspec_vsm.spots.StarSpot
"""
from astropy import units as u
import numpy as np
import pytest

from vspec_vsm.spots import StarSpot, SpotCollection, SpotGenerator
from vspec_vsm.coordinate_grid import RectangularGrid
from vspec_vsm.config import MSH

rng = np.random.default_rng(20)


def init_test_spot(**kwargs):
    """
    Initialize a default StarSpot for testing
    """
    lat = kwargs.get('lat', 0*u.deg)
    lon = kwargs.get('lon', 0*u.deg)
    Amax = kwargs.get('Amax', 500*MSH)
    A0 = kwargs.get('A0', 400*MSH)
    Teff_umbra = kwargs.get('Teff_umbra', 2700*u.K)
    Teff_penumbra = kwargs.get('Teff_penumbra', 2500*u.K)
    r_A = kwargs.get('r_A', 5.)
    growing = kwargs.get('growing', True)
    growth_rate = kwargs.get('growth_rate', 0/u.day)
    decay_rate = kwargs.get('decay_rate', 0*MSH/u.day)
    grid_params = kwargs.get('grid_params', 1000)
    gridmaker = kwargs.get('gridmaker', None)

    spot = StarSpot(
        lat=lat,
        lon=lon,
        area_max=Amax,
        area_current=A0,
        teff_umbra=Teff_umbra,
        teff_penumbra=Teff_penumbra,
        area_over_umbra_area=r_A,
        is_growing=growing,
        growth_rate=growth_rate,
        decay_rate=decay_rate,
        grid_params=grid_params,
        gridmaker=gridmaker
    )
    return spot


def test_spot_init():
    """
    Test __init__() for starspots
    """
    spot = init_test_spot(grid_params=20000)
    assert np.all(spot.r > 0*u.deg)
    assert np.all(spot.r < 180*u.deg)
    assert np.max(spot.r.to_value(u.deg)) == pytest.approx(180, abs=1)

    other = init_test_spot(
        gridmaker=RectangularGrid(500+1, 1000),
    )
    assert spot.gridmaker != other.gridmaker


def test_spot_str():
    """
    Make sure the __str__ method works.
    """
    spot = init_test_spot()
    s = str(spot)
    assert isinstance(s, str)


def test_spot_radius():
    """
    Test the spot radius
    """
    current_area = 400*MSH
    spot = init_test_spot(A0=current_area)
    # A = pi r**2
    # r = sqrt(A/pi)
    pred = np.sqrt(current_area/np.pi)
    obs = spot.radius
    assert pred.to_value(u.km) == pytest.approx(obs.to_value(u.km), rel=1e-6)


def test_spot_angular_radius():
    """
    Tests for the angular radius of the spot.
    """
    stellar_rad = 1000*u.km
    star_surface_area = 4*np.pi*stellar_rad**2
    spot = init_test_spot(A0=star_surface_area)

    pred = 180*u.deg
    obs = spot.angular_radius(stellar_rad)
    assert obs.to_value(u.deg) == pytest.approx(
        pred.to_value(u.deg), abs=0.01)

    spot = init_test_spot(A0=0.5*star_surface_area)
    pred = 90*u.deg
    obs = spot.angular_radius(stellar_rad)
    assert obs.to_value(u.deg) == pytest.approx(
        pred.to_value(u.deg), abs=0.01)


def test_spot_map_pixels():
    """
    Test StarSpot.map_pixels
    """
    stellar_rad = 1000*u.km
    star_surface_area = 4*np.pi*stellar_rad**2
    spot = init_test_spot(A0=star_surface_area, r_A=1)  # no penumbra
    pixmaps = spot.map_pixels(stellar_rad)
    umbra = pixmaps[spot.teff_umbra]
    penumbra = pixmaps[spot.teff_penumbra]
    assert np.all(umbra)
    assert not np.any(penumbra & ~umbra)

    spot = init_test_spot(A0=star_surface_area, r_A=np.inf)  # no umbra
    pixmaps = spot.map_pixels(stellar_rad)
    umbra = pixmaps[spot.teff_umbra]
    penumbra = pixmaps[spot.teff_penumbra]
    assert np.all(penumbra)
    assert not np.any(umbra)

    spot = init_test_spot(A0=star_surface_area, r_A=4, lat=0*u.deg)
    # half and half --> This is a sphere, so our 2D approximation goes out the window
    # for a large radius spot.
    pixmaps = spot.map_pixels(stellar_rad)
    umbra = pixmaps[spot.teff_umbra]
    penumbra = pixmaps[spot.teff_penumbra]
    lats, _ = spot.gridmaker.grid()
    sin_theta = np.cos(lats)
    assert (umbra*sin_theta.value).sum() == pytest.approx(((penumbra &
                                                            ~umbra)*sin_theta.value).sum(), rel=0.01)

    spot = init_test_spot(A0=1e-3*star_surface_area,
                          r_A=2, grid_params=(2000,4000))
    # half and half --> we try to get close to 2D without the discrete grid messing things up
    pixmaps = spot.map_pixels(stellar_rad)
    umbra = pixmaps[spot.teff_umbra]
    penumbra = pixmaps[spot.teff_penumbra]
    lats, _ = spot.gridmaker.grid()
    sin_theta = np.cos(lats)
    assert (umbra*sin_theta.value).sum() == pytest.approx(
        ((penumbra & ~umbra)*sin_theta.value).sum(), rel=0.05)


def test_spot_surface_fraction():
    """
    Test StarSpot.sufrace_fraction
    """
    stellar_rad = 1000*u.km
    star_surface_area = 4*np.pi*stellar_rad**2
    spot = init_test_spot(A0=0.01*star_surface_area)
    sub_obs = {'lat': 0*u.deg, 'lon': 0*u.deg}
    assert spot.surface_fraction(
        sub_obs, stellar_rad) == pytest.approx(0.02, rel=0.01)
    sub_obs = {'lat': 0*u.deg, 'lon': 180*u.deg}
    assert spot.surface_fraction(
        sub_obs, stellar_rad) == pytest.approx(0.0, rel=0.01)


def test_spot_age():
    """
    Test StarSpot.age
    """
    spot = init_test_spot(A0=10*MSH, Amax=100*MSH, growth_rate=1 /
                          u.day, decay_rate=10*MSH/u.day, growing=True)
    step = 1*u.day
    assert spot.area_current.to_value(MSH) == pytest.approx(10, rel=0.01)
    spot.age(step)
    assert spot.area_current.to_value(MSH) == pytest.approx(20, rel=0.01)
    spot.age(step)
    assert spot.area_current.to_value(MSH) == pytest.approx(40, rel=0.01)
    spot.age(step)
    assert spot.area_current.to_value(MSH) == pytest.approx(80, rel=0.01)
    spot.age(step)
    assert not spot.is_growing

    spot = init_test_spot(A0=100*MSH, Amax=100*MSH, growth_rate=1 /
                          u.day, decay_rate=10*MSH/u.day, growing=False)
    assert spot.area_current.to_value(MSH) == pytest.approx(100, rel=0.01)
    spot.age(step)
    assert spot.area_current.to_value(MSH) == pytest.approx(90, rel=0.01)
    spot.age(9*step)
    assert spot.area_current.to_value(MSH) == pytest.approx(0, rel=0.01)

    spot = init_test_spot(A0=10*MSH, Amax=100*MSH, growth_rate=1 /
                          u.day, decay_rate=10*MSH/u.day, growing=True)
    step = 1*u.day
    assert spot.area_current.to_value(MSH) == pytest.approx(10, rel=0.01)
    spot.age(20*step)
    assert spot.area_current.to_value(MSH) == pytest.approx(0, rel=0.01)


def test_init_spot_collection():
    """
    Test spot initialization
    """
    N = 4
    collec = SpotCollection(*[init_test_spot()
                            for i in range(N)], grid_params=(300, 600))
    expected_grid = RectangularGrid(300, 600)
    for spot in collec.spots:
        assert isinstance(spot, StarSpot)
    assert collec.gridmaker == expected_grid
    for spot in collec.spots:
        assert spot.gridmaker == collec.gridmaker


def test_spot_collection_add_spot():
    """
    Test `SpotCollection.add_spot()`
    """
    N = 4
    collec = SpotCollection(*[init_test_spot()
                            for i in range(N)], grid_params=(300, 600))
    assert len(collec.spots) == N
    new_spot = init_test_spot()
    collec.add_spot(new_spot)
    assert len(collec.spots) == N+1
    new_spots = [init_test_spot() for i in range(N)]
    collec.add_spot(new_spots)
    assert len(collec.spots) == 2*N+1
    for spot in collec.spots:
        assert spot.gridmaker == collec.gridmaker


def test_spot_collection_clean_spotlist():
    """
    Test `SpotCollection.clean_spotlist()`
    """
    N = 4
    spots = [init_test_spot() for i in range(N)] + \
        [init_test_spot(A0=0*MSH, growing=True)]
    collec = SpotCollection(*spots, grid_params=(300, 600))
    collec.clean_spotlist()
    assert len(collec.spots) == N+1
    spots = [init_test_spot() for i in range(N)] + \
        [init_test_spot(A0=0*MSH, growing=False)]
    collec = SpotCollection(*spots, grid_params=(300, 600))
    collec.clean_spotlist()
    assert len(collec.spots) == N


def test_spot_collection_map_pixels():
    """
    Test `SpotCollection.map_pixels()`
    """
    R_star = 0.15*u.R_sun
    Teff = 3300*u.K
    spots = [
        init_test_spot(Teff_umbra=2500*u.K, Teff_penumbra=2500*u.K),
        init_test_spot(Teff_umbra=2700*u.K, Teff_penumbra=2700*u.K)
    ]
    collec = SpotCollection(*spots, grid_params=(300, 600))
    pmap = collec.map_pixels(R_star, Teff)
    assert not np.any(pmap == 2700*u.K)

    spots = [
        init_test_spot(Teff_umbra=3500*u.K, Teff_penumbra=3500*u.K),
        init_test_spot(Teff_umbra=3700*u.K, Teff_penumbra=3700*u.K)
    ]
    collec = SpotCollection(*spots, grid_params=(300, 600))
    pmap = collec.map_pixels(R_star, Teff)
    assert np.all(pmap == Teff)


def test_spot_collection_age():
    """
    Test `SpotCollection.age()`
    """
    spot = init_test_spot(A0=10*MSH, Amax=100*MSH, growth_rate=1 /
                          u.day, decay_rate=10*MSH/u.day, growing=True)
    time = 1*u.day
    collec = SpotCollection(spot, grid_params=(300, 600))
    assert collec.spots[0].area_current == 10*MSH
    collec.age(time)
    assert collec.spots[0].area_current == 20*MSH
    collec.age(time*20)
    assert len(collec.spots) == 0


def init_spot_generator(**kwargs):
    """
    Initialize a `SpotGenerator` for testing
    """
    return SpotGenerator(
        dist_area_mean=kwargs.get('average_area', 100*MSH),
        dist_area_logsigma=kwargs.get('area_spread', 0.1),
        umbra_teff=kwargs.get('umbra_teff', 2500*u.K),
        penumbra_teff=kwargs.get('penumbra_teff', 2700*u.K),
        growth_rate=kwargs.get('growth_rate', 0/u.day),
        decay_rate=kwargs.get('decay_rate', 0*MSH/u.day),
        init_area=kwargs.get('starting_size', 10*MSH),
        distribution=kwargs.get('distribution', 'iso'),
        coverage=kwargs.get('coverage', 0.0),
        grid_params=kwargs.get('grid_params', (300, 600)),
        rng=rng
    )


def test_spot_generator_init():
    """
    Test `SpotGenerator.__init__()`
    """
    Nlat, Nlon = 300, 600
    gen = init_spot_generator(grid_params=(Nlat, Nlon))
    grid = RectangularGrid(Nlat, Nlon)
    assert gen.gridmaker == grid


def test_spot_generator_is_static():
    gen = init_spot_generator(growth_rate=0/u.day, decay_rate=0*MSH/u.day)
    assert gen.is_static
    gen = init_spot_generator(growth_rate=0/u.day, decay_rate=10*MSH/u.day)
    assert not gen.is_static


def test_spot_generator_mean_lifetime():
    growth_rate = 0.5/u.day
    decay_rate = 10*MSH/u.day
    max_size = 100*MSH
    starting_size = 10*MSH
    gen = init_spot_generator(growth_rate=growth_rate, decay_rate=decay_rate,
                              average_area=max_size, starting_size=starting_size)
    expected_value = max_size/decay_rate - \
        np.log(starting_size/max_size)/growth_rate
    assert expected_value == gen.mean_lifetime


def test_spot_generator_N_spots_to_birth():
    """
    Test `SpotGenerator.get_N_spots_to_birth()`
    """
    time = 10*u.day
    r_star = 1*u.R_sun
    gen = init_spot_generator(average_area=100*MSH,
                              coverage=0.0, decay_rate=10*MSH/u.day, growth_rate=0.5/u.day)
    exp = 0.0
    assert gen.get_n_spots_to_birth(
        time, r_star) == pytest.approx(exp, abs=1e-6)


def test_spot_generator_get_coordinates():
    """
    Test `SpotGenerator.get_coordinates()`
    """
    N = 3
    gen = init_spot_generator(distribution='iso')
    lats, lons = gen.get_coordinates(N)
    assert len(lats) == N
    assert lats.unit == u.deg
    assert len(lons) == N
    assert lons.unit == u.deg

    gen = init_spot_generator(distribution='even')
    with pytest.raises(ValueError):
        gen.get_coordinates(N)


def test_spot_generator_generate_spots():
    """
    Test `SpotGenerator.generate_spots()`
    """
    N = 3
    gen = init_spot_generator(starting_size=10*MSH)
    spots = gen.generate_spots(N)
    assert len(spots) == N
    for spot in spots:
        assert spot.area_current == 10*MSH


def test_spot_generator_generate_mature_spots():
    """
    Test `SpotGenerator.generate_mature_spots()`
    """
    r_star = 0.15*u.R_sun
    starting_size = 10*MSH
    mean_spot_area = 500*MSH
    gen = SpotGenerator(
        dist_area_mean=mean_spot_area,
        dist_area_logsigma=0.2,
        umbra_teff=2600*u.K,
        penumbra_teff=2900*u.K,
        growth_rate=0/u.day,
        decay_rate=0*MSH/u.day,
        init_area=starting_size,
        distribution='iso',
        coverage=0.2,
        grid_params=(500, 1000), rng=rng
    )
    with pytest.raises(ValueError):
        gen.generate_mature_spots(-0.1, r_star)
    with pytest.raises(ValueError):
        gen.generate_mature_spots(1.1, r_star)
    spots = gen.generate_mature_spots(0.1, r_star)
    assert np.any([spot.area_current > starting_size for spot in spots])

    gen = init_spot_generator(
        starting_size=starting_size, growth_rate=1/u.day, decay_rate=10*MSH/u.day)
    spots = gen.generate_mature_spots(0.1, r_star)
    assert np.any([spot.area_current > starting_size for spot in spots])
