# Copyright (c) 2019-2021 Thomas Kramer.
# SPDX-FileCopyrightText: 2022 Thomas Kramer
#
# SPDX-License-Identifier: GPL-3.0-or-later

from typing import Any, List, Dict, Optional, Tuple
from itertools import chain
from .boolean_functions import parse_boolean_function, format_boolean_function
from .arrays import strings_to_array, array_to_strings
import numpy as np
from sympy.logic import boolalg
import sympy


class Define:
    def __init__(self, attribute_name, group_name, attribute_type):
        """

        :param attribute_name: Name of the new defined attribute.
        :param group_name: Name of the group in which the attribute is created.
        :param attribute_type: Data type of the attribute: boolean, string, integer or float
        """

        self.attribute_name = attribute_name
        self.group_name = group_name
        self.attribute_type = attribute_type

    def __str__(self):
        return 'define ("{}", "{}", "{}")'.format(self.attribute_name, self.group_name, self.attribute_type)

    def __repr__(self):
        return str(self)


class Attribute:
    def __init__(self, attribute_name: str, value):
        """
        :param attribute_name: Name of the attribute.
        :param value: Value of the attribute.
        An attribute value can have types
        `float`, `string`, `EscapedString`, `ArithExpression`, `WithUnit`, `NameBitSelection`, `List[float]`.
        """

        self.name = attribute_name
        self.value = value

    def __str__(self):
        return '{} : {}'.format(self.name, self.value)

    def __repr__(self):
        return str(self)


class Group:
    def __init__(self, group_name: str,
                 args: List[str] = None,
                 attributes: List[Attribute] = None,
                 groups: List = None,
                 defines: List[Define] = None):
        self.groups: List[Group] = groups if groups is not None else []

        self.group_name = group_name
        self.args = args if args is not None else []
        assert isinstance(self.args, list)

        self.attributes = attributes if attributes is not None else list()
        assert isinstance(self.attributes, list)
        assert all((isinstance(v, Attribute) for v in self.attributes))

        self.groups = groups if groups is not None else []
        assert isinstance(self.groups, list)
        assert all((isinstance(v, Group) for v in self.groups))

        self.defines: List[Define] = defines if defines is not None else []
        assert isinstance(self.defines, list)
        assert all((isinstance(v, Define) for v in self.defines))

    def get_groups(self, type_name: str, argument: Optional[str] = None) -> List:
        """ Get all groups of type `type_name`.
        Optionally filter the groups by their first argument.
        :param type_name:
        :param argument:
        :return: List[Group]
        """
        return [g for g in self.groups
                if g.group_name == type_name
                and (argument is None or
                     (len(g.args) > 0 and g.args[0] == argument)
                     )
                ]

    def get_group(self, type_name: str, argument: Optional[str] = None):
        """
        Get exactly one group of type `type_name`.
        :param type_name:
        :return: Group
        """
        groups = self.get_groups(type_name, argument=argument)
        assert len(groups) == 1, "There must be exactly one instance of group '{}'. " \
                                 "Found {}.".format(type_name, len(groups))
        return groups[0]

    def __repr__(self) -> str:
        return "%s (%s) t{%s, %s}" % (self.group_name, self.args, self.attributes, self.groups)

    def __str__(self) -> str:
        """
        Create formatted string representation that can be dumped to a liberty file.
        :return:
        """
        return "\n".join(self._format())

    def _format(self, indent: str = " " * 2) -> List[str]:
        """
        Create the liberty file format line by line.
        :return: A list of lines.
        """

        def format_value(v) -> str:
            if v is None:
                return ''
            return str(v)

        define_lines = list()
        for d in self.defines:
            define_lines.append('{};'.format(d))

        sub_group_lines = [g._format(indent=indent) for g in self.groups]

        attr_lines = list()

        for attr in self.attributes:
            attr_name, attr_value = attr.name, attr.value
            if isinstance(attr_value, list):
                # Complex attribute
                formatted = [format_value(x) for x in attr_value]

                if any((isinstance(x, EscapedString) for x in attr_value)):
                    attr_lines.append('{} ( \\'.format(attr_name))
                    for i, l in enumerate(formatted):
                        if i < len(formatted) - 1:
                            end = ', \\'
                        else:
                            end = ' \\'
                        attr_lines.append(indent + l + end)
                    attr_lines.append(');')
                else:
                    values = "({})".format(", ".join(formatted))
                    attr_lines.append("{} {};".format(attr_name, values))
            else:
                # Simple attribute
                values = format_value(attr_value)
                attr_lines.append("{} : {};".format(attr_name, values))

        lines = list()
        lines.append("{} ({}) {{".format(self.group_name, ", ".join([format_value(f) for f in self.args])))

        for l in chain(define_lines, attr_lines, *sub_group_lines):
            lines.append(indent + l)

        lines.append("}")

        return lines

    def get_attributes(self, key: str) -> List[Any]:
        """
        Find attributes values by attribute name.
        :param key: The name of the attribute.
        :return: Returns a list of attribute values.
        """
        return [a.value for a in self.attributes if a.name == key]

    def get_attribute(self, key: str, default=None) -> Any:
        """
        Find exactly one attribute value based on its name.
        Raises an exception if there is no or more than one attributes with this name.
        :param key: Name of the attribute.
        :param default: Returns this default value if no attribute with this name is found.
        :return: The attribute value.
        """
        attrs = self.get_attributes(key)

        if len(attrs) == 0:
            return default

        assert len(attrs) == 1, "Expected to find exactly one attribute with name '{}'. " \
                                "Found {}.".format(key, len(attrs))
        return attrs[0]

    def __getitem__(self, item):
        return self.get_attribute(item)

    def __setitem__(self, key, value):
        """
        Replace the first occurrence of the attribute with the given value.
        """
        
        for a in self.attributes:
            if a.name == key:
                a.value = value
                return
            
        # Otherwise append a new attribute.
        self.attributes.append(Attribute(key, value))

    def __delitem__(self, key):
        """
        Delete all attributes with the name `key`.
        """
        self.attributes = [attr for attr in self.attributes if attr.name != key]

    def __contains__(self, item):
        return len(self.get_attributes(item)) > 0

    def get(self, key, default=None):
        return self.get_attribute(key, default)

    def get_array(self, key) -> np.ndarray:
        """
        Helper function for reading array falues.
        Get a 1D or 2D array as a numpy.ndarray object.
        :param key: Name of the attribute.
        :return: ndarray
        """
        str_array = self[key]
        str_array = [s.value for s in str_array]
        return strings_to_array(str_array)

    def set_array(self, key, value: np.ndarray):
        """
        Helper functions for storing array values.
        Set the value of the attribute "key" to the given array value.
        If attributes with the given name already exist, then the first of them will be replaced.
        """
        str_array = array_to_strings(value)
        str_array = [EscapedString(s) for s in str_array]
        self[key] = str_array

    def get_boolean_function(self, key) -> boolalg.Boolean:
        """
        Get parsed boolean expression.
        Intended for getting the value of the `function` attribute of pins.
        :param key:
        :return: Returns none if there is no function defined under this key.
        """

        f_str = self.get(key)
        if f_str is None:
            return None
        f = parse_boolean_function(f_str.value)
        return f

    def set_boolean_function(self, key, boolean: boolalg.Boolean):
        """
        Format the boolean expression and store it as an attribute with name `key`.
        :param key:
        :param boolean: Sympy boolean expression.
        """
        f_str = format_boolean_function(boolean)
        self[key] = '"{}"'.format(f_str)


class CellGroup(Group):

    def __init__(self, cell_name: str, attributes: List[Attribute],
                 sub_groups: List[Group]):
        super().__init__("cell", args=[cell_name], attributes=attributes,
                         groups=sub_groups)
        self.name = cell_name


class WithUnit:
    """
    Store a value with a unit attached.
    """

    def __init__(self, value, unit: str):
        self.value = value
        self.unit = unit

    def __str__(self):
        return "{}{}".format(self.value, self.unit)

    def __repr__(self):
        return str(self)

    def __eq__(self, other):
        if not isinstance(other, WithUnit):
            return False
        else:
            return self.value == other.value and self.unit == other.unit


class EscapedString:

    def __init__(self, value):
        self.value = value

    def __str__(self):
        return '"{}"'.format(self.value)

    def __repr__(self):
        return str(self)

    def __eq__(self, other):
        if isinstance(other, EscapedString):
            return self.value == other.value
        else:
            return self.value == other

    def __hash__(self):
        return hash(self.value)


class ArithExpression:
    """
    Arithmetic expression like `(VDD + 0.5) * 1.2`.
    Arithmetic expressions are stored as strings.
    """

    def __init__(self, value: str):
        self.value = value

    def __str__(self):
        return '"{}"'.format(self.value)

    def __repr__(self):
        return str(self)

    def to_sympy_expression(self) -> sympy.Expr:
        return sympy.parse_expr(self.value)


class NameBitSelection:
    """Name with bit selection (e.g. ADR[32:0])"""

    def __init__(self, name, sel1, sel2=None):
        self.name = name
        self.sel1 = sel1
        self.sel2 = sel2

    def __str__(self):
        if self.sel2 is not None:
            sel = "{}:{}".format(int(self.sel1), int(self.sel2))
        else:
            sel = int(self.sel1)
        return "{}[{}]".format(self.name, sel)

    def __repr__(self):
        return str(self)


def select_cell(library: Group, cell_name: str) -> Optional[Group]:
    """
    Select a cell by name from a library group.
    :param library:
    :param cell_name:
    :return:
    """
    available_cell_names = {g.args[0] for g in library.get_groups('cell')}

    if cell_name in available_cell_names:
        return library.get_group('cell', cell_name)
    else:
        raise KeyError("Cell name must be one of: {}".format(list(sorted(available_cell_names))))


def select_pin(cell: Group, pin_name: str) -> Optional[Group]:
    """
    Select a pin by name from a cell group.
    :param cell:
    :param pin_name:
    :return:
    """
    available_pin_names = {g.args[0] for g in cell.get_groups('pin')}

    if pin_name in available_pin_names:
        return cell.get_group('pin', pin_name)
    else:
        raise KeyError("Pin name must be one of: {}".format(list(sorted(available_pin_names))))


def select_timing_group(pin: Group,
                        related_pin: str,
                        when: Optional[str] = None,
                        timing_type: Optional[str] = None) -> Optional[Group]:
    """
    Select a timing group by name from a pin group.
    :param pin:
    :param related_pin:
    :param when:
    :param timing_type: Select by 'timing_type' attribute.
    :return:
    """

    assert isinstance(pin, Group)
    assert pin.group_name == 'pin', "`pin` must be a pin-group but is '{}'".format(pin.group_name)

    if when is not None:
        when = parse_boolean_function(when)
        
    def unwrap_escaped_string(s):
        if isinstance(s, EscapedString):
            return s.value
        return s
    
    # Select timing groups by value of the 'related_pin' attribute.
    timing_groups = [g
                     for g in pin.get_groups('timing')
                     if 'related_pin' in g
                     and unwrap_escaped_string(g['related_pin']) == related_pin
                     ]

    if not timing_groups:
        # Notify the user what `related_pin`s could have been chosen.
        raise KeyError(("No timing group found. Related pin name must be one of: {}".
            format(sorted(list({
            g['related_pin'].value
            for g in pin.get_groups('timing')
            if 'related_pin' in g
        })))
        ))

    # Select by `when`.
    if when is not None:

        def expr_equal(a, reference) -> str:
            """
            Test if the two expressions are equal.
            """
            a = unwrap_escaped_string(a)
            a = parse_boolean_function(a)
            return a == reference
            
        timing_groups_selected = [g
                         for g in timing_groups
                         if 'when' in g
                         and expr_equal(g['when'], when)
                         ]
        if not timing_groups_selected:
            # Notify the user what `related_pin`s could have been chosen.
            raise KeyError(("No timing group found. `when` must be one of: {}".
                format(sorted(list({
                unwrap_escaped_string(g['when'])
                for g in timing_groups
                if 'when' in g
            })))
            ))
        timing_groups = timing_groups_selected

    # Select by timing_type
    if timing_type is not None:
        timing_groups_selected = [g
                         for g in timing_groups
                         if 'timing_type' in g
                         and g['timing_type'] == timing_type
                         ]
        if not timing_groups_selected:
            # Notify the user what `timing_type`s could have been chosen.
            raise KeyError(("No timing group found. `timing_type` must be one of: {}".
                format(sorted(list({
                unwrap_escaped_string(g['timing_type'])
                for g in timing_groups
            })))
            ))
        timing_groups = timing_groups_selected

    timing_group = timing_groups[0]
    return timing_group


def select_timing_table(pin: Group,
                        related_pin: str,
                        table_name: str,
                        when: Optional[str] = None,
                        timing_type: Optional[str] = None) -> Optional[Group]:
    """
    Get a timing table by name from a pin group.
    :param pin: The pin group of the output pin.
    :param related_pin: The related input pin name.
    :param table_name: The name of the timing table. ('cell_rise', 'cell_fall', 'rise_transition', ...)
    :param timing_type: Select by 'timing_type' attribute.
    :return:
    """

    # Find timing group.
    timing_group = select_timing_group(pin, related_pin, timing_type=timing_type, when=when)

    available_table_names = {g.group_name for g in timing_group.groups}

    if table_name in available_table_names:
        return timing_group.get_group(table_name)
    else:
        raise KeyError(("Table name must be one of: {}".format(list(sorted(available_table_names)))))
