"""Unit tests for the DMD class.
"""

# third party packages
from pytest import raises
import torch as pt
# flowtorch packages
from flowtorch import DATASETS
from flowtorch.data import FOAMDataloader
from .dmd import _dft_properties, DMD


def test_DMD():
    path = DATASETS["of_cavity_binary"]
    loader = FOAMDataloader(path)
    times = loader.write_times[1:]
    data = loader.load_snapshot("p", times)
    rows, cols = data.shape
    rank = 3
    dmd = DMD(data, dt=0.1, rank=rank)
    assert dmd.eigvals.shape == (rank,)
    assert dmd.eigvals.dtype == pt.complex64
    assert dmd.eigvecs.shape == (rank, rank)
    assert dmd.eigvecs.dtype == pt.complex64
    assert dmd.modes.dtype == pt.complex64
    assert dmd.modes.shape == (rows, rank)
    assert dmd.frequency.shape == (rank,)
    assert dmd.growth_rate.shape == (rank,)
    assert dmd.amplitude.shape == (rank,)
    assert dmd.amplitude.dtype == pt.complex64
    assert dmd.dynamics.shape == (rank, cols)
    assert dmd.dynamics.dtype == pt.complex64
    assert dmd.integral_contribution.shape == (rank,)
    assert dmd.integral_contribution.dtype == pt.float32
    assert dmd.reconstruction.shape == (rows, cols)
    assert dmd.reconstruction.dtype == data.dtype
    partial = dmd.partial_reconstruction({0})
    assert partial.dtype == data.dtype
    assert partial.shape == (rows, cols)
    partial = dmd.partial_reconstruction({0, 2})
    assert partial.dtype == data.dtype
    assert partial.shape == (rows, cols)
    top = dmd.top_modes(10)
    top = dmd.top_modes(10, True)
    assert top.shape == (min(rank, 10),)
    assert top.dtype == pt.int64
    assert dmd.reconstruction_error.shape == (rows, cols)
    assert dmd.projection_error.shape == (rows, cols - 1)
    dft = dmd.dft_properties
    assert len(dft) == 3
    assert dft == (10.0, 5.0, 10.0/(cols - 1))
    # robust DMD
    dmd = DMD(data, dt=0.1, rank=rank, robust=True)
    assert dmd.svd.L.shape == (data.shape[0], rank+1)
    # unitary operator
    dmd = DMD(data, dt=0.1, rank=rank, unitary=True)
    operator = dmd.operator
    shape = operator.shape
    assert shape == (rank, rank)
    diag = operator.conj().T @ operator
    assert pt.allclose(diag, pt.diag(pt.ones(rank)), atol=1e-6)
    # optimal mode amplitudes
    dmd = DMD(data, dt=0.1, rank=rank, optimal=True)
    dmd = DMD(data, dt=0.1, rank=rank, unitary=True, optimal=True)
    assert dmd.amplitude.shape == (rank,)
    assert dmd.amplitude.dtype == pt.complex64
    # total least-squares
    dmd = DMD(data, dt=0.1, tlsq=True)
    assert dmd.amplitude.dtype == pt.complex64
    dmd = DMD(data, dt=0.1, rank=rank, optimal=True, tlsq=True)
    assert dmd.amplitude.shape == (rank,)
    DX, DY = dmd.tlsq_error
    assert  DX.shape == (data.shape[0], data.shape[1] - 1)
    assert  DY.shape == (data.shape[0], data.shape[1] - 1)
    # predict member function
    dmd = DMD(data, dt=0.1, rank=3)
    prediction = dmd.predict(data[:, -1], 10)
    assert prediction.shape == (rows, 11)
    assert prediction.dtype == data.dtype
    data = data.type(pt.complex64)
    dmd = DMD(data, dt=0.1, rank=3)
    prediction = dmd.predict(data[:, -1], 10)
    assert prediction.shape == (rows, 11)
    assert prediction.dtype == data.dtype


def test_dft_properties():
    props = _dft_properties(0.1, 500)
    assert len(props) == 3
    assert props == (10.0, 5.0, 10.0/500)






