"""Common strategies"""
from typing import cast

import numpy as np
from hypothesis import strategies as hst
from hypothesis.extra import numpy as hnp
from hypothesis.strategies import composite, DrawFn, SearchStrategy
from numpy._typing import NDArray

from zema_emc_annotated.data_types import UncertainArray


@composite
def uncertain_arrays(
    draw: DrawFn,
    greater_than: float = -1e2,
    less_than: float = 1e2,
    samples: int | None = None,
    scaler: int | None = None,
) -> SearchStrategy[UncertainArray]:
    if samples is None:
        samples = draw(hst.integers(min_value=1, max_value=10))
    if scaler is None:
        scaler = draw(hst.integers(min_value=1, max_value=10))
    values: NDArray[np.float64] = cast(
        NDArray[np.float64],
        draw(
            hnp.arrays(
                dtype=np.float64,
                shape=(samples, scaler * 11),
                elements=hst.floats(min_value=greater_than, max_value=less_than),
            )
        ),
    )
    std_uncertainties = cast(
        NDArray[np.float64],
        draw(
            hnp.arrays(
                dtype=np.float64,
                shape=values.shape,
                elements=hst.floats(
                    min_value=np.abs(values).min() * 1e-3,
                    max_value=np.abs(values).min() * 1e2,
                ),
            )
        ),
    )
    return cast(
        SearchStrategy[UncertainArray],
        UncertainArray(values, std_uncertainties),
    )