# -*- coding: utf-8 -*-
from os import remove
from os.path import dirname

import numpy as np
import pytest

from pysme.iliffe_vector import Iliffe_vector
from pysme.sme import SME_Structure as SME_Struct


@pytest.fixture
def cwd():
    return dirname(__file__)


@pytest.fixture
def filename(cwd):
    fname = "{}/__test.sme".format(cwd)
    yield fname
    try:
        remove(fname)
    except:
        pass


def test_empty_structure():
    """Test that all properties behave well when nothing is set"""
    empty = SME_Struct()

    assert isinstance(empty.version, str)
    assert empty.teff is not None
    assert empty.logg is not None
    assert empty.vmic == 0
    assert empty.vmac == 0
    assert empty.vsini == 0

    assert empty.nseg == 0
    assert empty.wave is None
    assert empty.spec is None
    assert empty.uncs is None
    assert empty.synth is None
    assert empty.cont is None
    assert empty.mask is None
    assert empty.mask_good is None
    assert empty.mask_bad is None
    # assert empty.mask_line is None
    # assert empty.mask_continuum is None

    assert empty.cscale.shape == (0, 1)
    assert empty.vrad.shape == (0,)
    assert empty.cscale_flag == "none"
    assert empty.vrad_flag == "none"
    assert empty.cscale_degree == 0

    assert empty.mu is not None
    assert empty.nmu == 7

    # assert empty.md5 is not None

    assert empty.linelist is not None
    assert empty.species is not None
    assert len(empty.species) == 0
    assert empty.atomic is not None

    assert empty.monh == 0
    assert not np.isnan(empty["abund Fe"])
    assert empty.abund["H"] == 12
    assert not np.isnan(empty.abund()["Mg"])

    assert empty.system_info is not None
    assert empty.system_info.arch == ""

    assert len(empty.fitparameters) == 0
    assert empty.fitresults is not None
    assert empty.fitresults.covariance is None

    assert empty.atmo is not None
    assert empty.atmo.depth is None

    assert empty.nlte is not None
    assert empty.nlte.elements == []


def test_save_and_load_structure(filename):
    sme = SME_Struct()
    assert sme.teff is not None

    sme.teff = 5000
    sme.save(filename)
    del sme
    sme = SME_Struct.load(filename)
    assert sme.teff == 5000

    remove(filename)

    data = np.linspace(1000, 2000, 100)
    sme.wave = data
    sme.spec = data
    sme.save(filename)
    sme = SME_Struct.load(filename)
    assert np.all(sme.wave[0] == data)
    assert np.all(sme.spec[0] == data)
    assert sme.nseg == 1


def test_load_idl_savefile(cwd):
    filename = "{}/testcase1.inp".format(cwd)
    sme = SME_Struct.load(filename)

    assert sme.teff == 5770
    assert sme.wave is not None

    assert sme.nseg == 1
    assert sme.cscale_flag == "linear"
    assert sme.vrad_flag == "each"


def test_cscale_degree():
    sme = SME_Struct()
    sme.cscale = 1

    flags = ["none", "fix", "constant", "linear", "quadratic"]
    degrees = [0, 0, 0, 1, 2]

    for f, d in zip(flags, degrees):
        sme.cscale_flag = f
        assert sme.cscale_degree == d
        assert sme.cscale.shape[0] == 0
        assert sme.cscale.shape[1] == d + 1


def test_idlver():
    sme = SME_Struct()
    sme.system_info.update()
    # assert sme.idlver.arch == "x86_64"


def test_fitresults():
    sme = SME_Struct()
    sme.fitresults.chisq = 100
    sme.fitresults.clear()
    assert sme.fitresults.chisq is None


def test_wint_single_segment_array_input():
    sme = SME_Struct()
    wint = np.linspace(5000.0, 5001.0, 11)

    sme.wint = wint

    assert isinstance(sme.wint, Iliffe_vector)
    assert sme.wint.nseg == 1
    assert np.allclose(sme.wint[0], wint)


def test_wint_multi_segment_list_input():
    sme = SME_Struct()
    wint0 = np.linspace(5000.0, 5001.0, 5)
    wint1 = np.linspace(6000.0, 6002.0, 7)

    sme.wint = [wint0, wint1]

    assert isinstance(sme.wint, Iliffe_vector)
    assert sme.wint.nseg == 2
    assert np.allclose(sme.wint[0], wint0)
    assert np.allclose(sme.wint[1], wint1)


def test_wave_reassignment_is_not_silently_truncated():
    sme = SME_Struct()
    short = np.arange(3700.0, 3700.012, 0.002)
    long = np.arange(3700.0, 3970.505, 0.0002)

    sme.wave = short
    sme.wave = long

    assert sme.wave.nseg == 1
    assert len(sme.wave[0]) == len(long)
    assert np.allclose(sme.wave[0], long)


def test_wave_reassignment_invalidates_wave_dependent_state():
    sme = SME_Struct()
    old_wave = np.linspace(5000.0, 5001.0, 11)
    new_wave = np.linspace(6000.0, 6002.0, 21)

    sme.wave = old_wave
    sme.spec = old_wave.copy()
    sme.uncs = np.ones_like(old_wave)
    sme.mask = np.ones_like(old_wave, dtype=np.int32)
    sme.synth = np.ones_like(old_wave) * 0.9
    sme.cont = np.ones_like(old_wave)
    sme.telluric = np.ones_like(old_wave)
    sme.sint = [np.ones((3, old_wave.size))]
    sme.cint = [np.ones((3, old_wave.size))]

    sme.vrad_flag = "fix"
    sme.vrad = np.array([12.0])
    sme.vrad_unc = np.array([[0.1, 0.2]])
    sme.cscale_flag = "linear"
    sme.cscale = np.array([[0.0, 1.0]])
    sme.cscale_unc = np.zeros((1, 2, 2))

    sme.wave = new_wave

    assert len(sme.wave[0]) == len(new_wave)
    assert sme.spec is None
    assert sme.uncs is None
    assert sme.mask is None
    assert sme.synth is None
    assert sme.cont is None
    assert sme.telluric is None
    assert sme.sint is None
    assert sme.cint is None

    # vrad/cscale are reset to defaults for the new segment layout.
    assert np.allclose(sme.vrad, [0.0])
    assert np.allclose(sme.cscale, [[0.0, 1.0]])
    assert sme.vrad_unc is None
    assert sme.cscale_unc is None
