"""
A module containing unit tests for the `tpwcs` module.

Licensed under a 3-clause BSD style license - see LICENSE.rst

"""
import copy
import pytest

import numpy as np
from astropy.modeling.models import Scale, Identity
from astropy import wcs as fitswcs
from astropy.modeling import CompoundModel
import gwcs
from gwcs.geometry import SphericalToCartesian
from tweakwcs.linearfit import build_fit_matrix
from tweakwcs.correctors import JWSTWCSCorrector, FITSWCSCorrector

from .helper_correctors import (make_mock_jwst_wcs, make_mock_jwst_pipeline,
                                DummyWCSCorrector, create_DetToV2V3,
                                create_V2V3ToDet)


_ATOL = 100 * np.finfo(np.array([1.]).dtype).eps


def test_tpwcs():
    tpwcs = DummyWCSCorrector(None, meta={})
    x, y, ra, dec = np.random.random(4)
    matrix = np.random.random((2, 2))
    shift = np.random.random(2)

    assert tpwcs.world_to_det(ra, dec) == (ra, dec)
    assert tpwcs.world_to_tanp(ra, dec) == (ra, dec)
    assert tpwcs.det_to_world(x, y) == (x, y)
    assert tpwcs.det_to_tanp(x, y) == (x, y)
    assert tpwcs.tanp_to_det(x, y) == (x, y)
    assert tpwcs.tanp_to_world(x, y) == (x, y)
    assert tpwcs.tanp_center_pixel_scale == 1
    assert tpwcs.wcs is None
    assert tpwcs.original_wcs is None
    assert isinstance(tpwcs.copy(), DummyWCSCorrector)
    assert tpwcs.bounding_box is None

    tpwcs.set_correction(matrix=matrix, shift=shift,
                         meta={'pytest': 'ABC.WCSCorrector'}, pytest_kwarg=True)
    assert np.all(tpwcs.meta['matrix'] == matrix)
    assert np.all(tpwcs.meta['shift'] == shift)
    assert tpwcs.meta['pytest'] == 'ABC.WCSCorrector'

    with pytest.raises(TypeError) as arg_err:
        tpwcs.set_correction(matrix, shift, None, {'pytest': 'ABC.WCSCorrector'},
                             'some_weird_arg')
    assert (
        arg_err.value.args[0].endswith(
            "set_correction() takes from 1 to 5 positional arguments but 6 were given"
        )
    )


def test_mock_jwst_gwcs():
    w = make_mock_jwst_wcs(v2ref=123, v3ref=500, roll=115, crpix=[-512, -512],
                           cd=[[1e-5, 0], [0, 1e-5]], crval=[82, 12])

    assert np.allclose(w.invert(*w(23, 1023)), (23, 1023))


@pytest.mark.parametrize('crpix, cd', [
    (np.zeros(3), np.diag(np.ones(3))),
    (np.zeros((2, 2)), np.diag(np.ones(2))),
])
def test_mock_wcs_fails(crpix, cd):
    from astropy.modeling import InputParameterError
    with pytest.raises(InputParameterError):
        make_mock_jwst_wcs(v2ref=123, v3ref=500, roll=15, crpix=crpix,
                           cd=cd, crval=[82, 12])
    with pytest.raises(InputParameterError):
        create_DetToV2V3(v2ref=123, v3ref=500, roll=15, crpix=crpix, cd=cd)
    with pytest.raises(InputParameterError):
        create_V2V3ToDet(v2ref=123, v3ref=500, roll=15, crpix=crpix, cd=cd)


def test_v2v3todet_roundtrips():
    s2c = (Scale(1.0 / 3600.0) & Scale(1.0 / 3600.0)) | SphericalToCartesian(wrap_lon_at=180)

    s = 1.0e-5
    crpix = np.random.random(2)
    alpha = 0.25 * np.pi * np.random.random()
    x, y = 1024 * np.random.random(2)
    v2, v3 = 45 * np.random.random(2)
    cd = [[s * np.cos(alpha), -s * np.sin(alpha)],
          [s * np.sin(alpha), s * np.cos(alpha)]]

    d2v = create_DetToV2V3(
        v2ref=123.0, v3ref=500.0, roll=15.0, crpix=crpix, cd=cd
    )
    v2d = create_V2V3ToDet(
        v2ref=123.0, v3ref=500.0, roll=15.0, crpix=crpix, cd=cd
    )

    assert np.allclose(d2v.inverse(*d2v(x, y)), (x, y),
                       rtol=1e3 * _ATOL, atol=1e3 * _ATOL)

    assert (
        np.allclose(
            s2c(*v2d.inverse(*v2d(v2, v3))),
            s2c(v2, v3),
            rtol=1e3 * _ATOL, atol=_ATOL
        ) or np.allclose(
            -np.asanyarray(s2c(*v2d.inverse(*v2d(v2, v3)))),
            s2c(v2, v3),
            rtol=1e3 * _ATOL, atol=_ATOL
        )
    )
    assert np.allclose(v2d(*d2v(x, y)), (x, y),
                       rtol=1e5 * _ATOL, atol=1e3 * _ATOL)

    assert (
        np.allclose(
            s2c(*d2v(*v2d(v2, v3))),
            s2c(v2, v3),
            rtol=1e3 * _ATOL, atol=1e3 * _ATOL
        ) or np.allclose(
            -np.asanyarray(s2c(*d2v(*v2d(v2, v3)))),
            s2c(v2, v3),
            rtol=1e3 * _ATOL, atol=1e3 * _ATOL
        )
    )


def test_jwst_wcs_corr_applied(mock_jwst_wcs):
    w = make_mock_jwst_wcs(
        v2ref=123.0, v3ref=500.0, roll=115.0, crpix=[512.0, 512.0],
        cd=[[1.0e-5, 0.0], [0.0, 1.0e-5]], crval=[82.0, 12.0]
    )

    wc = JWSTWCSCorrector(
        w, {'v2_ref': 123.0, 'v3_ref': 500.0, 'roll_ref': 115.0}, meta={}
    )
    wc.set_correction(meta={'dummy_meta': None}, dummy_par=1)
    assert 'v2v3corr' in wc.wcs.available_frames
    assert 'dummy_meta' in wc.meta


def test_jwst_wcs_corr_are_being_combined(mock_jwst_wcs):
    wc = JWSTWCSCorrector(
        mock_jwst_wcs, {'v2_ref': 123.0, 'v3_ref': 500.0, 'roll_ref': 115.0}
    )
    matrix1 = np.array([[1.0, 0.2], [-0.3, 1.1]])
    shift1 = np.array([5.0, -7.0])
    wc.set_correction(matrix=matrix1, shift=shift1)
    assert 'v2v3corr' in wc.wcs.available_frames

    matrix2 = np.linalg.inv(matrix1)
    shift2 = -np.dot(matrix2, shift1)
    wc = JWSTWCSCorrector(
        wc.wcs, {'v2_ref': 123.0, 'v3_ref': 500.0, 'roll_ref': 115.0}
    )
    wc.set_correction(matrix=matrix2, shift=shift2)

    v2v3idx = [k for k, n in enumerate(wc.wcs.available_frames)
               if n == 'v2v3corr']

    assert len(v2v3idx) == 1

    tp_corr = wc.wcs.pipeline[v2v3idx[0] - 1].transform

    assert isinstance(tp_corr, CompoundModel)
    assert np.max(np.abs(tp_corr['tp_affine'].matrix.value - np.identity(2))) < _ATOL
    assert np.max(np.abs(tp_corr['tp_affine'].translation.value)) < _ATOL


def test_jwstgwcs_unsupported_wcs():
    dummy_wcs = gwcs.wcs.WCS(Identity(2), 'det', 'world')
    with pytest.raises(ValueError):
        JWSTWCSCorrector(dummy_wcs, {})


def test_jwstgwcs_inconsistent_ref(mock_jwst_wcs):
    wc = JWSTWCSCorrector(
        mock_jwst_wcs, {'v2_ref': 123.0, 'v3_ref': 500.0, 'roll_ref': 115.0},
    )
    wc.set_correction()

    with pytest.raises(ValueError):
        wc = JWSTWCSCorrector(
            wc.wcs, {'v2_ref': 124.0, 'v3_ref': 500.0, 'roll_ref': 115.0},
        )


def test_jwstgwcs_wrong_tpcorr_type(mock_jwst_wcs):
    wc = JWSTWCSCorrector(
        mock_jwst_wcs, {'v2_ref': 123.0, 'v3_ref': 500.0, 'roll_ref': 115.0},
    )
    wc.set_correction()
    p = wc.wcs.pipeline

    np = [(v.frame, create_V2V3ToDet()) if v.frame.name == 'v2v3vacorr'
          else v for v in p]
    mangled_wc = gwcs.wcs.WCS(np)

    with pytest.raises(ValueError):
        wc = JWSTWCSCorrector(
            mangled_wc, {'v2_ref': 123.0, 'v3_ref': 500.0, 'roll_ref': 115.0},
        )


def test_jwstgwcs_ref_angles_preserved(mock_jwst_wcs):
    wc = JWSTWCSCorrector(
        mock_jwst_wcs, {'v2_ref': 123.0, 'v3_ref': 500.0, 'roll_ref': 115.0},
    )
    assert wc.ref_angles['v2_ref'] == 123.0
    assert wc.ref_angles['v3_ref'] == 500.0
    assert wc.ref_angles['roll_ref'] == 115.0

# Test inputs with different shapes


@pytest.mark.parametrize('inputs', [[500, 512, 12, 24],
                                    [[500, 500], [512, 512], [12, 12], [24, 24]],
                                    [[[500, 500], [500, 500]],
                                     [[512, 512], [512, 512]],
                                     [[12, 12], [12, 12]], [[24, 24], [24, 24]]]
                                    ])
def test_jwstgwcs_detector_to_world(inputs):
    w = make_mock_jwst_wcs(
        v2ref=0.0, v3ref=0.0, roll=0.0, crpix=[500.0, 512.0],
        cd=[[1.0e-5, 0.0], [0.0, 1.0e-5]], crval=[12.0, 24.0]
    )
    wc = JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})
    wc.set_correction()
    x, y, ra, dec = inputs
    assert np.allclose(wc.det_to_world(x, y), (ra, dec), atol=_ATOL)
    assert np.allclose(wc.world_to_det(ra, dec), (x, y), atol=_ATOL)


@pytest.mark.parametrize('inputs', [[0, 0, 12, 24],
                                    [[0, 0], [0, 0], [12, 12], [24, 24]],
                                    [[[0, 0], [0, 0]],
                                     [[0, 0], [0, 0]],
                                     [[12, 12], [12, 12]], [[24, 24], [24, 24]]]
                                    ])
def test_jwstgwcs_tangent_to_world(inputs):
    w = make_mock_jwst_wcs(
        v2ref=0.0, v3ref=0.0, roll=0.0, crpix=[500.0, 512.0],
        cd=[[1.0e-5, 0.0], [0.0, 1.0e-5]], crval=[12.0, 24.0]
    )
    wc = JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})
    wc.set_correction()
    tanp_x, tanp_y, ra, dec = inputs
    assert np.allclose(wc.world_to_tanp(ra, dec), (tanp_x, tanp_y), atol=1000 * _ATOL)
    assert np.allclose(wc.tanp_to_world(tanp_x, tanp_y), (ra, dec), atol=_ATOL)


@pytest.mark.parametrize('inputs', [[500, 512, 0, 0],
                                    [[500, 500], [512, 512], [0, 0], [0, 0]],
                                    [[[500, 500], [500, 500]],
                                     [[512, 512], [512, 512]],
                                     [[0, 0], [0, 0]], [[0, 0], [0, 0]]]
                                    ])
def test_jwstgwcs_detector_to_tanp(inputs):
    w = make_mock_jwst_wcs(
        v2ref=0.0, v3ref=0.0, roll=0.0, crpix=[500.0, 512.0],
        cd=[[1.0e-5, 0.0], [0.0, 1.0e-5]], crval=[12.0, 24.0]
    )
    wc = JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})
    wc.set_correction()
    x, y, tanp_x, tanp_y = inputs
    assert np.allclose(wc.det_to_tanp(x, y), (tanp_x, tanp_y), atol=10 * _ATOL)
    assert np.allclose(wc.tanp_to_det(tanp_x, tanp_y), (x, y), atol=_ATOL)


def test_jwstgwcs_bbox():
    w = make_mock_jwst_wcs(
        v2ref=0.0, v3ref=0.0, roll=0.0, crpix=[500.0, 512.0],
        cd=[[1.0e-5, 0.0], [0.0, 1.0e-5]], crval=[12.0, 24.0]
    )
    wc = JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})
    wc.set_correction()

    assert np.allclose(
        wc.bounding_box,
        ((-0.5, 1024 - 0.5), (-0.5, 2048 - 0.5)),
        atol=_ATOL
    )

    wc._owcs.bounding_box = None
    assert np.allclose(
        wc.bounding_box,
        ((-0.5, 1024 - 0.5), (-0.5, 2048 - 0.5)),
        atol=_ATOL
    )

    wc._owcs.array_shape = None
    assert wc.bounding_box is None


def test_jwstgwcs_bad_pipelines_no_vacorr():
    p0 = make_mock_jwst_pipeline(
        v2ref=0.0, v3ref=0.0, roll=0.0, crpix=[500.0, 512.0],
        cd=[[1.0e-5, 0.0], [0.0, 1.0e-5]], crval=[12.0, 24.0],
        enable_vacorr=False
    )

    # no pipeline or empty pipeline:
    with pytest.raises(ValueError):
        JWSTWCSCorrector(None, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})

    # fewer than 3 frames:
    w = gwcs.wcs.WCS(p0[:2])
    with pytest.raises(ValueError):
        JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})

    # repeated (any one of the) last two frames:
    w = gwcs.wcs.WCS(p0 + [p0[-1]])
    with pytest.raises(ValueError):
        JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})

    w = gwcs.wcs.WCS(p0 + [p0[-2]])
    with pytest.raises(ValueError):
        JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})

    # multiple 'v2v3' frames:
    w = gwcs.wcs.WCS(p0)
    w = JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})
    p = w.wcs.pipeline
    p.insert(1, p[1])
    w = gwcs.wcs.WCS(p)
    with pytest.raises(ValueError):
        JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})

    # misplaced 'v2v3' frame:
    w = gwcs.wcs.WCS(p0)
    w = JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})
    w.set_correction()
    p = w.wcs.pipeline
    del p[0]
    w = gwcs.wcs.WCS(p)
    with pytest.raises(ValueError):
        JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})

    # multiple 'v2v3corr' frame:
    w = gwcs.wcs.WCS(p0)
    w = JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})
    w.set_correction()
    p = w.wcs.pipeline
    p.insert(1, p[-2])
    w = gwcs.wcs.WCS(p)
    with pytest.raises(ValueError):
        JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})

    # misplaced 'v2v3corr' frame:
    del p[-2]
    w = gwcs.wcs.WCS(p)
    with pytest.raises(ValueError):
        JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})


def test_jwstgwcs_bad_pipelines_with_vacorr():
    p0 = make_mock_jwst_pipeline(
        v2ref=0.0, v3ref=0.0, roll=0.0, crpix=[500.0, 512.0],
        cd=[[1.0e-5, 0.0], [0.0, 1.0e-5]], crval=[12.0, 24.0],
        enable_vacorr=True
    )

    # multiple 'v2v3vacorr' frames:
    w = gwcs.wcs.WCS(p0)
    w = JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})
    p = w.wcs.pipeline
    p.insert(1, p[2])
    w = gwcs.wcs.WCS(p)
    with pytest.raises(ValueError):
        JWSTWCSCorrector(w, {'v2_ref': 0.0, 'v3_ref': 0.0, 'roll_ref': 0.0})


def test_fitswcs_non_celestial():
    # non-celestial WCS
    w = fitswcs.WCS(naxis=3)
    with pytest.raises(ValueError):
        FITSWCSCorrector(w)

    # invalid WCS
    with pytest.raises(ValueError):
        FITSWCSCorrector(None)


def test_fitswcs_unaccounted_dist(mock_fits_wcs):
    w = copy.deepcopy(mock_fits_wcs)
    w.pix2foc = lambda x, o: x + 3
    with pytest.raises(ValueError):
        FITSWCSCorrector(w)

    w = copy.deepcopy(mock_fits_wcs)
    f = w.all_world2pix
    w.all_world2pix = lambda x, o: f(x, o) + 2
    with pytest.raises(ValueError):
        FITSWCSCorrector(w)


def test_fitswcs_1(mock_fits_wcs):
    wc = FITSWCSCorrector(mock_fits_wcs)
    wc.set_correction()


def test_fitswcs_coord_transforms(mock_fits_wcs):
    w = fitswcs.WCS(naxis=2)
    w.wcs.cd = build_fit_matrix(0, 1e-5)
    w.wcs.crval = [12, 24]
    w.wcs.crpix = [500, 512]
    w.wcs.ctype = ['RA---TAN', 'DEC--TAN']
    w.pixel_shape = [1024, 2048]
    w.wcs.set()

    wc = FITSWCSCorrector(w)
    wc.set_correction()

    assert np.allclose(wc.det_to_world(499, 511), (12, 24), atol=_ATOL)
    assert np.allclose(wc.world_to_det(12, 24), (499, 511), atol=_ATOL)
    assert np.allclose(wc.det_to_tanp(499, 511), (499, 511), atol=_ATOL)
    assert np.allclose(wc.tanp_to_det(499, 511), (499, 511), atol=_ATOL)
    assert np.allclose(wc.world_to_tanp(12, 24), (499, 511), atol=1e-8)
    assert np.allclose(wc.tanp_to_world(499, 511), (12, 24), atol=_ATOL)


def test_fitswcs_bbox(mock_fits_wcs):
    w = copy.deepcopy(mock_fits_wcs)
    wc = FITSWCSCorrector(w)
    wc.set_correction()

    assert np.allclose(
        wc.bounding_box,
        ((-0.5, 1024 - 0.5), (-0.5, 2048 - 0.5)),
        atol=_ATOL
    )

    wc._owcs.pixel_bounds = None
    assert np.allclose(
        wc.bounding_box,
        ((-0.5, 1024 - 0.5), (-0.5, 2048 - 0.5)),
        atol=_ATOL
    )

    wc._owcs.bounding_box = None
    assert np.allclose(
        wc.bounding_box,
        ((-0.5, 1024 - 0.5), (-0.5, 2048 - 0.5)),
        atol=_ATOL
    )

    wc._owcs.array_shape = None
    assert wc.bounding_box is None
