"""AframeXR components"""

import copy
import json
import marimo
from typing import Literal

from aframexr.api.data import Data, URLData
from aframexr.api.encoding import *
from aframexr.api.filters import *
from aframexr.utils.constants import *
from aframexr.utils.scene_creator import SceneCreator


class TopLevelMixin:
    """Top level chart class."""

    def __init__(self, specs: dict = None):
        if specs is None:  # For calls of Chart.__init__(), calling super().__init__()
            self._specifications = {}  # Specifications of the scene, in JSON format
        else:  # For calls of __add__(), to create a new object and do not modify the rest
            self._specifications = specs

    # Concatenating charts
    def __add__(self, other):
        """
        Concatenation of charts (place charts in the same scene).
        Creates and returns a new scene with the charts. The original charts are not modified.
        """

        if not isinstance(other, TopLevelMixin):
            raise TypeError(f"Cannot add {type(other).__name__} to {type(self).__name__}")

        # Create the new concatenation specifications
        concat_specs = {'concat': []}

        # Look if there is concatenation of concatenated charts (to join all in one concat field)
        if self._specifications.get('concat'):
            for chart in self._specifications['concat']:
                concat_specs['concat'].append(chart)
        else:
            concat_specs['concat'].append(self._specifications)

        if other._specifications.get('concat'):
            for chart in other._specifications['concat']:
                concat_specs['concat'].append(chart)
        else:
            concat_specs['concat'].append(other._specifications)

        concatenated_chart = TopLevelMixin(concat_specs)  # Create a new 'scene' to preserve the original charts
        return concatenated_chart

    # Copy of the chart
    def copy(self):
        """Returns a deep copy of the chart."""

        return copy.deepcopy(self)

    # Importing charts
    @staticmethod
    def from_dict(specs: dict) -> 'TopLevelMixin':
        """
        Import the chart from the JSON dict specifications.

        Parameters
        ----------
        specs : dict
            JSON specifications of the chart.

        Raises
        ------
        TypeError
            If specs is not a dictionary.
        """

        if not isinstance(specs, dict):
            raise TypeError(f'Expected dict, got {type(specs).__name__} instead.')
        return TopLevelMixin(specs)

    @staticmethod
    def from_json(specs: str) -> 'TopLevelMixin':
        """
        Create the chart from the JSON string specifications.

        Parameters
        ----------
        specs : str
            JSON specifications of the chart.

        Raises
        ------
        TypeError
            If specs is not a string.
        """

        if not isinstance(specs, str):
            raise TypeError(f'Expected str, got {type(specs).__name__} instead.')
        return TopLevelMixin(json.loads(specs))

    # Exporting charts
    def save(self, fp: str, fileFormat: Literal['json', 'html'] = None):
        """
        Saves the chart into a file, supported formats are JSON and HTML.

        Parameters
        ----------
        fp : str
            File path.
        fileFormat : str (optional)
            Format of the file could be ['html', 'json'].
            If no format is specified, the chart will be saved depending on the file extension.

        Raises
        ------
        NotImplementedError
            If fileFormat is invalid.
        """

        if fileFormat == 'html' or fp.endswith('.html'):
            with open(fp, 'w') as file:
                file.write(self.to_html())
        elif fileFormat == 'json' or fp.endswith('.json'):
            with open(fp, 'w') as file:
                json.dump(self._specifications, file, indent=4)
        else:
            raise NotImplementedError('That format is not supported.')

    # Showing the scene
    def show(self):
        """Show the scene in the Marimo notebook."""

        html_scene = SceneCreator.create_scene(self._specifications)
        return marimo.iframe(html_scene)

    # Chart formats
    def to_dict(self) -> dict:
        """Returns the scene specifications as a dictionary."""

        return self._specifications

    def to_html(self) -> str:
        """Returns the HTML representation of the scene."""

        return SceneCreator.create_scene(self._specifications)

    def to_json(self) -> str:
        """Returns the JSON string of the scene."""

        return json.dumps(self._specifications)


class Chart(TopLevelMixin):
    """
    Simple chart class.

    Parameters
    ----------
    data : Data | URLData
        Data or URLData object of the data.

    position : str (optional)
        Position of the chart. The format is: 'x y z'. Refers to the position for the origin of coordinate system.
    rotation : str (optional)
        Rotation of the chart in degrees. The format is: 'x y z'. The rotation axis is the coordinate system.

    Raises
    ------
    TypeError
        If data is not a Data or URLData object.
    ValueError
        If position or rotation are invalid.
    """

    def __init__(self, data: Data | URLData, position: str = DEFAULT_CHART_POS, rotation: str = DEFAULT_CHART_ROTATION):
        super().__init__()

        # Data
        if isinstance(data, Data):
            self._specifications.update({'data': {'values': data.values}})
        elif isinstance(data, URLData):
            self._specifications.update({'data': {'url': data.url}})
        else:
            raise TypeError(f'Expected Data | URLData, got {type(data).__name__} instead.')

        # Position
        pos_axes = position.strip().split()
        if len(pos_axes) != 3:
            raise ValueError(f'The position: {position} is not correct. Must be "x y z".')
        for axis in pos_axes:
            try:
                float(axis)
            except ValueError:
                raise ValueError(f'The position values must be numeric.')
        self._specifications.update({'position': f'{pos_axes[0]} {pos_axes[1]} {pos_axes[2]}'})

        # Rotation
        rot_axes = rotation.strip().split()
        if len(rot_axes) != 3:
            raise ValueError(f'The rotation: {rotation} is not correct. Must be "x y z".')
        for axis in rot_axes:
            try:
                float(axis)
            except ValueError:
                raise ValueError(f'The rotation values must be numeric.')
        self._specifications.update({'rotation': f'{rot_axes[0]} {rot_axes[1]} {rot_axes[2]}'})

    # Types of charts
    def mark_arc(self, radius: float = DEFAULT_PIE_RADIUS):
        """
        Pie chart and doughnut chart.

        Parameters
        ----------
        radius : float (optional)
            Outer radius of the pie chart. If not specified, using default. Must be greater than 0.
        """

        self._specifications.update({'mark': {'type': 'arc'}})
        if radius > 0:
            self._specifications['mark'].update({'radius': radius})
        else:
            raise ValueError('radius must be greater than 0.')
        return self

    def mark_bar(self, size: float = DEFAULT_BAR_WIDTH, height: float = DEFAULT_MAX_HEIGHT):
        """
        Bars chart.

        Parameters
        ----------
        size : float (optional)
            Width of the bars. If not specified, using default. Must be greater than 0.
        height : float (optional)
            Maximum height of the chart (the highest bar). If not specified, using default. Must be greater than 0.
        """

        self._specifications.update({'mark': {'type': 'bar'}})
        if size > 0:
            self._specifications['mark'].update({'width': size})
        else:
            raise ValueError('size must be greater than 0.')
        if height > 0:
            self._specifications.update({'height': height})
        else:
            raise ValueError('height must be greater than 0.')
        return self

    def mark_gltf(self, scale: str = DEFAULT_GLTF_SCALE):
        """
        GLTF model.

        Parameters
        ----------
        scale : str (optional)
            Scale of the GLTF model (from its original scale).

            **Format: 'x y z'** (values can be negative, works like a mirror).

            If an axis value is not specified, that value will be 1 (for example, '2 2' is the same as '2 2 1').
        """

        self._specifications.update({'mark': {'type': 'gltf'}})
        self._specifications['mark'].update({'scale': scale})
        return self

    def mark_image(self, width: int = DEFAULT_IMAGE_WIDTH, height: int = DEFAULT_IMAGE_HEIGHT):
        """
        Image.

        Parameters
        ----------
        width : int (optional)
            Width of the image. If not specified, using default. Must be greater than 0.
        height : int (optional)
            Height of the image. If not specified, using default. Must be greater than 0.

        Raises
        ------
        ValueError
            If width or height are not greater than 0.
        """

        self._specifications.update({'mark': {'type': 'image'}})
        if width > 0:
            self._specifications['mark'].update({'width': width})
        else:
            raise ValueError('width must be greater than 0.')
        if height > 0:
            self._specifications['mark'].update({'height': height})
        else:
            raise ValueError('height must be greater than 0.')
        return self

    def mark_point(self, size: float = DEFAULT_POINT_RADIUS, height: float = DEFAULT_MAX_HEIGHT):
        """
        Scatter plot and bubble chart.

        Parameters
        ----------
        size : float (optional)
            Maximum radius of the point. If not specified, using default. Must be greater than 0.
        height : float (optional)
            Maximum height of the chart. If not specified, using default. Must be greater than 0.

        Raises
        ------
        ValueError
            If size or height are not greater than 0.
        """

        self._specifications.update({'mark': {'type': 'point'}})
        if size > 0:
            self._specifications['mark'].update({'max_radius': size})
        else:
            raise ValueError('size must be greater than 0.')
        if height > 0:
            self._specifications.update({'height': height})
        else:
            raise ValueError('height must be greater than 0.')
        return self

    # Parameters of the chart
    def encode(self, color: str = '', size: str = '', theta: str = '', x: str | X = '', y: str | Y = '',
               z: str | Z = ''):
        """
        Add properties to the chart.

        Parameters
        ----------
        color : str (optional)
            Field of the data that will determine the color of the sphere in the scatter plot.
        size : str (optional)
            Field of the data that will determine the size of the sphere in the bubble chart (must be quantitative).
        theta : str (optional)
            Field of the data that will determine the arcs of the pie and doughnut chart (must be quantitative).
        x : str | X (optional)
            Field of the data that will determine the x-axis of the chart.
        y : str | Y (optional)
            Field of the data what will determine the y-axis of the chart.
        z : str | Z (optional)
            Field of the data what will determine the z-axis of the chart.

        Raises
        ------
        TypeError
            If the encoding type is incorrect.
        ValueError
            If the encoding values are incorrect.
        """

        filled_params = {}  # Dictionary that will store the parameters that have been filled

        # Verify the type of the arguments and store the filled parameters
        if color:
            if not isinstance(color, str):
                raise TypeError(f'Expected color as str, got {type(color).__name__} instead.')
            filled_params.update({'color': color})
        if size:
            if not isinstance(size, str):
                raise TypeError(f'Expected size as str, got {type(size).__name__} instead.')
            filled_params.update({'size': size})
        if theta:
            if not isinstance(theta, str):
                raise TypeError(f'Expected theta as str, got {type(theta).__name__} instead.')
            filled_params.update({'theta': theta})
        if x:
            if not isinstance(x, str | X):
                raise TypeError(f'Expected x as str | aframexr.X, got {type(x).__name__} instead.')
            filled_params.update({'x': x})
        if y:
            if not isinstance(y, str | Y):
                raise TypeError(f'Expected y as str | aframexr.Y, got {type(y).__name__} instead.')
            filled_params.update({'y': y})
        if z:
            if not isinstance(z, str | Z):
                raise TypeError(f'Expected z as str | aframe.Z, got {type(z).__name__} instead.')
            filled_params.update({'z': z})

        # Verify the argument combinations
        if self._specifications['mark']['type'] in ['bar', 'point'] and sum([x != '', y != '', z != '']) < 2:
            raise ValueError('at least 2 of (x, y, z) must be specified.')
        if self._specifications['mark']['type'] == 'arc' and (not theta or not color):
            if not theta: raise ValueError('theta must be specified in arc chart.')
            if not color: raise ValueError('color must be specified in arc chart.')
        if self._specifications['mark']['type'] == 'bar' and (color or size):
            if color: raise ValueError('bar chart does not support color.')
            if size: raise ValueError('bar chart does not support size.')

        # Do the encoding
        self._specifications.update({'encoding': {}})
        for param_key in filled_params:
            param_value = filled_params[param_key]
            if isinstance(param_value, Encoding):
                self._specifications['encoding'].update(param_value.to_dict())
            else:
                self._specifications['encoding'].update({param_key: {'field': param_value}})
        return self

    # Filtering data
    def transform_filter(self, equation_filter: str | FilterTransform):
        """
        Filters the chart with the given transformation.

        Parameters
        ----------
        equation_filter : str | FilterTransform
            The equation string of the filter transformation, or a Filter object (see Examples).

        Raises
        ------
        TypeError
            If equation is not a string or a Filter object.

        Notes
        -----
        Can be concatenated with the rest of functions of the Chart, without needing an asignation.

        Examples
        --------
        *Using transform_filter() giving the equation string:*

        >>> import aframexr
        >>> data = aframexr.URLData('./data.json')
        >>> filtered_chart = aframexr.Chart(data).mark_bar().encode(x='model', y='sales')
        >>> filtered_chart = filtered_chart.transform_filter('datum.motor=diesel')
        >>> #filtered_chart.show()

        *Using transform_filter() giving a Filter object*

        >>> import aframexr
        >>> data = aframexr.URLData('./data.json')
        >>> filtered_chart = aframexr.Chart(data).mark_bar().encode(x='model', y='sales')
        >>> filter_object = aframexr.FieldEqualPredicate(field='motor', equal='diesel')
        >>> filtered_chart = filtered_chart.transform_filter(filter_object)
        >>> #filtered_chart.show()
        """

        # Create a copy of the chart (in case of assignation, to preserve the main chart)
        filt_chart = self.copy()

        # Validate the type of equation_filter and get a filter object from the equation_filter
        if isinstance(equation_filter, str):
            filter_transform = FilterTransform.from_string(equation_filter)
        elif isinstance(equation_filter, FilterTransform):
            filter_transform = equation_filter
        else:
            raise TypeError(f'Expected string or FilterTransform object, got {type(equation_filter).__name__}.')

        # Add the information of the filter object to the specifications
        if not filt_chart._specifications.get('transform'):  # First time filtering the chart
            filt_chart._specifications.update({'transform': [filter_transform.equation_to_dict()]})  # Create field
        else:  # Not the first filter of the chart
            filt_chart._specifications['transform'].append(filter_transform.equation_to_dict())  # Add filter to field
        return filt_chart  # Returns the copy of the chart
