"""AframeXR axis creator"""

import polars as pl

from polars import Series
from typing import Literal

from .constants import *


_X_AXIS_LABELS_ROTATION = '-90 0 -90'
_Y_AXIS_LABELS_ROTATION = '0 0 0'
_Z_AXIS_LABELS_ROTATION = '-90 0 0'


def _get_labels_coords_for_quantitative_axis(axis_data: Series, axis_size: float) -> Series:
    """Returns the coordinates for the labels of the quantitative axis."""
    unique_values = axis_data.n_unique()

    if unique_values == 1:  # All the values are the same
        return Series([axis_size / 2])  # Only one tick is placed in the axis

    num_samples = unique_values if axis_data.dtype == pl.String else DEFAULT_NUM_OF_TICKS_IF_QUANTITATIVE_AXIS
    return pl.linear_space(  # Equally spaced values
            start=START_LABEL_OFFSET,  # Offset for the lowest label (for not being on the ground)
            end=axis_size,
            num_samples=num_samples,
            eager=True  # Returns a Series
        )

def _get_labels_values_for_quantitative_axis(axis_data: Series) -> Series:
    """Returns the values for the labels of the quantitative axis."""
    if axis_data.dtype == pl.String:  # Axis data contains nominal values, but user wants to encode as quantitative
        return axis_data.unique(maintain_order=True)  # Return the same values

    min_value, max_value = axis_data.min(), axis_data.max()

    if max_value == min_value:  # All the values are the same
        labels_values = Series([axis_data[0]])  # Only one tick is placed in the axis
    else:
        if max_value < 0:  # All data is negative
            start = min_value
            end = 0
        elif min_value > 0:  # All data is positive
            start = 0
            end = max_value
        else:
            start, end = min_value, max_value

        labels_values = pl.linear_space(
            start=start,
            end=end,
            num_samples=DEFAULT_NUM_OF_TICKS_IF_QUANTITATIVE_AXIS,
            eager=True  # Returns a Series
        )
    return labels_values


class AxisCreator:
    """Axis creator class."""

    @staticmethod
    def create_axis_html(start: str, end: str) -> str:
        """
        Create a line for the axis and returns its HTML.

        Parameters
        ----------
        start : str
            The base position of each axis.
        end : str
            The end position of the axis.
        """
        return f'<a-entity line="start: {start}; end: {end}; color: black"></a-entity>'

    @staticmethod
    def create_axis_specs(axis: Literal['x', 'y', 'z'], axis_data: Series, axis_encoding: str, axis_size: float,
                          elements_coords: Series, x_offset: float, y_offset: float, z_offset: float) -> dict:
        """Returns the axis specifications for x, y or z axis depending on its encoding."""
        axis_specs = {'start': None, 'end': None, 'labels_pos': [], 'labels_values': [], 'labels_rotation': '',
                      'labels_align': None}

        if axis_encoding == 'quantitative':
            coords = _get_labels_coords_for_quantitative_axis(axis_data, axis_size)
            labels_values = _get_labels_values_for_quantitative_axis(axis_data)
        elif axis_encoding == 'nominal':
            coords = elements_coords.unique(maintain_order=True)  # Align labels with elements
            labels_values = axis_data.unique(maintain_order=True)
        else:  # pragma: no cover (Encoding type must have been checked before)
            raise RuntimeError(f'Unreachable code. Check encoding type: {axis_encoding}')

        axis_specs['start'] = f'{x_offset} {y_offset} {z_offset}'
        if axis == 'x':
            axis_specs['end'] = f'{axis_size} {y_offset} {z_offset}'
            axis_specs['labels_pos'] = (coords.cast(pl.String) + f' {LABELS_Y_DELTA} {X_LABELS_Z_DELTA}').to_list()
            axis_specs['labels_rotation'] = _X_AXIS_LABELS_ROTATION
            axis_specs['labels_align'] = 'left'
        elif axis == 'y':
            axis_specs['end'] = f'{x_offset} {axis_size} {z_offset}'
            axis_specs['labels_pos'] = (f'{LABELS_X_DELTA} ' + coords.cast(pl.String) + ' 0').to_list()
            axis_specs['labels_rotation'] = _Y_AXIS_LABELS_ROTATION
            axis_specs['labels_align'] = 'right'
        elif axis == 'z':
            axis_specs['end'] = f'{x_offset} {y_offset} {-axis_size}'  # Negative axis size to go deep
            axis_specs['labels_pos'] = (f'{LABELS_X_DELTA} {LABELS_Y_DELTA} ' + coords.cast(pl.String)).to_list()
            axis_specs['labels_rotation'] = _Z_AXIS_LABELS_ROTATION
            axis_specs['labels_align'] = 'right'
        else:  # pragma: no cover (this method is only called by inner code methods; should be OK)
            raise RuntimeError('Unreachable code. Axis must be x or y or z')

        axis_specs['labels_values'] = labels_values.to_list()
        return axis_specs
