"""
Tests for codebook.per_round_max_decode method
"""

import numpy as np
import pandas as pd
import pytest

from starfish import Codebook, IntensityTable
from starfish.types import Features, Indices, SpotAttributes


def intensity_table_factory(data: np.ndarray=np.array([[[0, 3], [4, 0]]])) -> IntensityTable:
    """IntensityTable with a single feature that was measured over 2 channels and 2 rounds."""

    # generates spot attributes equal in size to the number of passed features.
    # each attribute has coordinates (z, y, x) equal to the feature index, and radius 1.
    spot_attributes_data = pd.DataFrame(
        data=np.array([[i, i, i, 1] for i in np.arange(data.shape[0])]),
        columns=[Indices.Z, Indices.Y, Indices.X, Features.SPOT_RADIUS]
    )

    spot_attributes = SpotAttributes(spot_attributes_data)
    intensity_table = IntensityTable.from_spot_data(data, spot_attributes)
    return intensity_table


def codebook_factory() -> Codebook:
    """
    Codebook with two codewords describing an experiment with three channels and two imaging rounds.
    Both codes have two "on" channels.
    """
    codebook_array = [
        {
            Features.CODEWORD: [
                {Indices.ROUND.value: 0, Indices.CH.value: 1, Features.CODE_VALUE: 1},
                {Indices.ROUND.value: 1, Indices.CH.value: 0, Features.CODE_VALUE: 1}
            ],
            Features.TARGET: 'GENE_A'
        },
        {
            Features.CODEWORD: [
                {Indices.ROUND.value: 0, Indices.CH.value: 1, Features.CODE_VALUE: 1},
                {Indices.ROUND.value: 1, Indices.CH.value: 1, Features.CODE_VALUE: 1}
            ],
            Features.TARGET: 'GENE_B'
        },
    ]
    return Codebook.from_code_array(codebook_array)


def test_intensity_tables_with_different_nubmers_of_codes_or_channels_throw_value_error():
    """
    The test passes a 3-round and 1-round IntensityTable to a 2-round codebook. Both should
    raise a ValueError.
    """
    data = np.array(
        [[[4, 3, 1],
          [4, 0, 2]]]
    )
    codebook = codebook_factory()
    intensities = intensity_table_factory(data)
    with pytest.raises(ValueError):
        codebook.decode_per_round_max(intensities)

    with pytest.raises(ValueError):
        codebook.decode_per_round_max(intensities.sel(round=1))


def test_example_intensity_decodes_to_gene_a():
    """The single intensity should decode to gene a. Verify that it does."""
    intensities = intensity_table_factory()
    codebook = codebook_factory()

    decoded_intensities = codebook.decode_per_round_max(intensities)

    assert decoded_intensities[Features.TARGET] == 'GENE_A'
    assert np.array_equal(decoded_intensities[Features.PASSES_THRESHOLDS], [True])
    assert np.array_equal(decoded_intensities[Features.DISTANCE], [0])


def test_example_intensity_that_has_no_match_is_assigned_nan():
    """
    The decoder assigns nan when it doesn't receive a match, but the nan is converted to a
    string due to the object dtype
    """
    no_match = np.array([[[3, 0], [0, 4]]])
    intensities = intensity_table_factory(no_match)
    codebook = codebook_factory()

    decoded_intensities = codebook.decode_per_round_max(intensities)

    # even though the feature did not decode, it should still be returned
    assert decoded_intensities.sizes[Features.AXIS] == 1

    # assert that the feature did not decode
    assert decoded_intensities[Features.TARGET] == 'nan'
    assert np.array_equal(decoded_intensities[Features.PASSES_THRESHOLDS], [False])

    # distance is calculated as the fraction of signal NOT in the selected channel. Here all of
    # the signal is in the max channel, so distance is 0.
    # Admittedly, this is a little bit confusing, since distance here measures both how correct
    # a code is, but also how wrong one is when it _doesn't_ decode properly.
    assert np.array_equal(decoded_intensities[Features.DISTANCE], [0])


def test_argmax_selects_the_last_equal_intensity_channel_and_decodes_consistently():
    """
    When two channels are tied, argmax selects the last first one. This is the code for Gene A:

    [[0, 1],
     [1, 0]]

    both of these observations could potentially match the code for GENE_A:

    [[4, 3],
     [4, 0]]

    and

    [[0, 3],
     [4, 3]]

    The first is interpreted as argmax as:

    [[1, 1],
     [0, 0]]

    which fails to match, while the second is interpreted as

    [[0, 1],
     [1, 0]]

    and matches GENE_A.

    This test also verifies that the decoder can operate on multiple codes.
    """

    data = np.array(
        [[[4, 3],  # this code is decoded "wrong"
          [4, 0]],
         [[0, 3],  # this code is decoded "right"
          [4, 3]]]
    )
    intensities = intensity_table_factory(data)
    codebook = codebook_factory()

    decoded_intensities = codebook.decode_per_round_max(intensities)
    assert np.array_equal(decoded_intensities[Features.TARGET].values, ['nan', 'GENE_A'])
