"""
Tests for fileio module
"""

# author: Andrew R. McCluskey (arm61)

import unittest
from datetime import datetime
import os.path

import pytest
import yaml
from orsopy.fileio.orso import Orso, OrsoDataset
from orsopy.fileio.data_source import (DataSource, Experiment, Sample,
                                       Measurement, InstrumentSettings)
from orsopy.fileio.reduction import Reduction, Software
from orsopy.fileio.base import Person, ValueRange, Value, File, Column
from orsopy.fileio.base import _validate_header_data, _read_header_data
from orsopy import fileio as fileio
import numpy as np


class TestOrso(unittest.TestCase):
    """
    Testing the Orso class.
    """
    def test_creation(self):
        """
        Creation of Orso object.
        """
        e = Experiment(
            'Experiment 1', 'ESTIA', datetime(2021, 7, 7, 16, 31, 10),
            'neutrons'
        )
        s = Sample('The sample')
        inst = InstrumentSettings(
            Value(4.0, 'deg'), ValueRange(2., 12., 'angstrom')
        )
        df = [File('README.rst', None)]
        m = Measurement(inst, df, scheme="angle-dispersive")
        p = Person('A Person', 'Some Uni')
        ds = DataSource(p, e, s, m)

        soft = Software('orsopy', '0.0.1', 'macOS-10.15')
        p2 = Person('Andrew McCluskey', 'European Spallation Source')
        redn = Reduction(
            soft, datetime(2021, 7, 14, 10, 10, 10),
            p2, ['footprint', 'background']
        )

        cols = [Column("Qz", unit='1/angstrom'), Column("R")]
        value = Orso(ds, redn, cols, 0)

        ds = value.data_source
        dsm = ds.measurement
        assert ds.owner.name == 'A Person'
        assert dsm.data_files[0].file == 'README.rst'
        assert dsm.instrument_settings.incident_angle.magnitude == 4.0
        assert dsm.instrument_settings.wavelength.min == 2.0
        assert dsm.instrument_settings.wavelength.max == 12.0
        assert value.reduction.software.name == 'orsopy'
        assert value.reduction.software.version == "0.0.1"
        assert value.reduction.timestamp == datetime(2021, 7, 14, 10, 10, 10)
        assert value.columns[0].name == 'Qz'
        assert value.columns[1].name == 'R'
        assert value.data_set == 0

        h = value.to_yaml()
        h = "\n".join(
            ["# ORSO reflectivity data file | 0.1 standard | YAML encoding"
             " | https://www.reflectometry.org/",
             h]
        )
        g = yaml.safe_load_all(h)
        _validate_header_data([next(g)])

    def test_creation_data_set1(self):
        """
        Creation of Orso object with a non-zero data_set.
        """
        e = Experiment(
            'Experiment 1', 'ESTIA', datetime(2021, 7, 7, 16, 31, 10),
            'neutrons'
        )
        s = Sample('The sample')
        inst = InstrumentSettings(
            Value(4.0, 'deg'), ValueRange(2., 12., 'angstrom')
        )
        df = [File('README.rst', None)]
        m = Measurement(inst, df, scheme="angle-dispersive")
        p = Person('A Person', 'Some Uni')
        ds = DataSource(p, e, s, m)

        soft = Software('orsopy', '0.0.1', 'macOS-10.15')
        p2 = Person('Andrew McCluskey', 'European Spallation Source')
        redn = Reduction(
            soft, datetime(2021, 7, 14, 10, 10, 10), p2,
            ['footprint', 'background']
        )

        cols = [Column("Qz", unit='1/angstrom'), Column("R")]
        value = Orso(ds, redn, cols, 1)

        dsm = value.data_source.measurement
        assert value.data_source.owner.name == 'A Person'
        assert dsm.data_files[0].file == 'README.rst'
        assert value.reduction.software.name == 'orsopy'
        assert value.columns[0].name == 'Qz'
        assert value.data_set == 1

        # check that data_set can also be a string.
        value = Orso(ds, redn, cols, 'fokdoks')
        assert value.data_set == 'fokdoks'
        # don't want class construction coercing a str to an int
        value = Orso(ds, redn, cols, '1')
        assert value.data_set == '1'

    def test_write_read(self):
        # test write and read of multiple datasets
        info = fileio.Orso.empty()
        info2 = fileio.Orso.empty()
        data = np.zeros((100, 3))
        data[:] = np.arange(100.0)[:, None]

        info.columns = [
            fileio.Column("Qz", "1/angstrom"),
            fileio.Column("R"),
            fileio.Column("sR"),
        ]
        info2.columns = info.columns
        info.data_source.measurement.instrument_settings.polarization = "p"
        info2.data_source.measurement.instrument_settings.polarization = "m"
        info.data_set = "up polarization"
        info2.data_set = "down polarization"
        info2.data_source.sample.comment = "this is a comment"

        ds = fileio.OrsoDataset(info, data)
        ds2 = fileio.OrsoDataset(info2, data)

        info3 = fileio.Orso(
            data_source=fileio.DataSource(
                sample=fileio.Sample(
                    name="My Sample",
                    type="solid",
                    description="Something descriptive",
                ),
                experiment=fileio.Experiment(
                    title="Main experiment",
                    instrument="Reflectometer",
                    date=datetime.now(),
                    probe="x-rays",
                ),
                owner=fileio.Person("someone", "important"),
                measurement=fileio.Measurement(
                    instrument_settings=fileio.InstrumentSettings(
                        incident_angle=fileio.Value(13.4, "deg"),
                        wavelength=fileio.Value(5.34, "A"),
                    ),
                    data_files=["abc", "def", "ghi"],
                    references=["more", "files"],
                    scheme="angle-dispersive",
                ),
            ),
            reduction=fileio.Reduction(software="awesome orso"),
            data_set="Filled header",
            columns=info.columns,
        )
        ds3 = fileio.OrsoDataset(info3, data)
        fileio.save_orso([ds, ds2, ds3], "test.ort", comment='Interdiffusion')

        ls1, ls2, ls3 = fileio.load_orso("test.ort")
        assert ls1 == ds
        assert ls2 == ds2
        assert ls3 == ds3

        _read_header_data("test.ort", validate=True)

    def test_unique_dataset(self):
        # checks that data_set is unique on saving of OrsoDatasets
        info = Orso.empty()
        info.data_set = 0
        info.columns = [Column("stuff")] * 4

        info2 = Orso.empty()
        info2.data_set = 0
        info2.columns = [Column("stuff")] * 4

        ds = OrsoDataset(info, np.empty((2, 4)))
        ds2 = OrsoDataset(info2, np.empty((2, 4)))

        with pytest.raises(ValueError):
            fileio.save_orso([ds, ds2], 'test_data_set.ort')

    def test_extra_elements(self):
        # if there are extra elements present in the ORT file they should still
        # be loadable. They won't be there as dataclass fields, but they'll be
        # visible as attributes.
        datasets = fileio.load_orso(os.path.join("tests", "test_example.ort"))
        info = datasets[0].info
        assert hasattr(
            info.data_source.measurement.instrument_settings.incident_angle,
            'resolution'
        )


class TestFunctions(unittest.TestCase):
    """
    Tests for functionality in the Orso module.
    """
    def test_make_empty(self):
        """
        Creation of the empty Orso object.
        """
        empty = Orso.empty()
        assert issubclass(empty.__class__, Orso)
        ds = empty.data_source
        assert ds.owner.name is None
        assert ds.experiment.title is None
        assert ds.experiment.instrument is None
        assert ds.experiment.date is None
        assert ds.experiment.probe is None
        assert ds.sample.name is None
        assert ds.measurement.instrument_settings.incident_angle.magnitude is None
        assert ds.measurement.instrument_settings.wavelength.magnitude is None
        assert ds.measurement.data_files is None
        assert empty.reduction.software.name is None
        assert empty.reduction.software.version is None
        assert empty.reduction.software.platform is None
        assert empty.reduction.timestamp is None
        assert empty.reduction.creator is None
        assert ds.owner.affiliation is None
        assert ds.sample.name is None
        assert empty.reduction.corrections is None
        assert empty.reduction.creator is None
        assert empty.columns == [Column.empty()]
        assert empty.data_set is None
        dct = empty.to_dict()
        _validate_header_data([dct])

    def test_empty_to_yaml(self):
        """
        Checking yaml string form empty Orso object.

        TODO: Fix once correct format is known.
        """
        empty = Orso.empty()
        req = (
            'data_source:\n  owner:\n    name: null\n'
            '    affiliation: null\n  experiment:\n    title: null\n'
            '    instrument: null\n    date: null\n    probe: null\n'
            '  sample:\n    name: null\n  measurement:\n'
            '    instrument_settings:\n      incident_angle:\n        magnitude: null\n'
            '      wavelength:\n        magnitude: null\n      polarization: unpolarized\n'
            '    data_files: null\nreduction:\n  software:\n    name: null\n'
            'columns:\n- name: null\n'
        )
        assert empty.to_yaml() == req
