import numpy as np

from neurosym.dsl.dsl import DSL
from neurosym.dsl.dsl_factory import DSLFactory


def example_dataset(num_sequences, len_sequences, *, seed, slack=30, stride=5):
    """
    An example dataset for a simple symbolic regression task. The sequences here
    are generated by taking an IID uniform distribution on ```[-2, 2]``` of length ```(len_sequences * stride + slack)```,
    then averaging over the slack window (compensating for variance by multiplying by sqrt(slack)),
    then taking every stride-th element.

    This is done in order to make the sequence more smooth and thus more "naturalistic"; however
    there is no guarantee that the sequences are generated by a simple function.

    :param num_sequences: The number of sequences to generate
    :param len_sequences: The length of each sequence
    :param seed: The random seed to use
    :param slack: The slack window to average over. Larger values will make the sequences more smooth
    :param stride: The stride to take when downsampling. Larger values will make the sequences more variant

    :return: A tuple of ```(xs, values)```, where xs is the x-axis and values is a matrix of shape ```(num_sequences, len_sequences)```
    """
    xs = np.linspace(-10, 10, len_sequences * stride)
    values = (
        np.random.RandomState(seed).rand(num_sequences, len_sequences * stride + slack)
        * 4
        - 2
    )
    values = np.mean(
        [values[:, i : i + len_sequences * stride] for i in range(slack)], axis=0
    )
    values *= np.sqrt(slack)
    values = values[:, ::stride]
    xs = xs[::stride]
    return xs, values


def example_dsl() -> DSL:
    """
    An example DSL for the simple symbolic regression task. This DSL includes
    basic arithmetic operations, as well as some trigonometric functions.
    """
    dslf = DSLFactory()
    dslf.production("0", "() -> f", lambda: 0)
    dslf.production("1", "() -> f", lambda: 1)
    dslf.production("2", "() -> f", lambda: 2)
    dslf.production("+", "(f, f) -> f", lambda x, y: x + y)
    dslf.production("-", "(f, f) -> f", lambda x, y: x - y)
    dslf.production("*", "(f, f) -> f", lambda x, y: x * y)
    dslf.production("**", "(f, f) -> f", lambda x, y: x**y)
    dslf.production("/", "(f, f) -> f", lambda x, y: x / y)
    dslf.production("sin", "f -> f", np.sin)
    dslf.production("sqrt", "f -> f", np.sqrt)
    dslf.production("<", "(f, f) -> b", lambda x, y: x < y)
    dslf.production("ite", "(b, f, f) -> f", lambda cond, x, y: x if cond else y)
    dslf.lambdas()
    dslf.prune_to("f -> f")
    return dslf.finalize()
