from typing import Tuple, Any, List, Dict

import matplotlib.pyplot as plt

from isthmuslib import VectorMultiset
from src.isthmuslib.logging import parse_string_with_embedded_json, parse_string_with_manual_tokens
from src.isthmuslib.utils import risky_cast, process_queue, recursive_batch_evaluation, return_best_input


def test_risky_cast():
    def val_and_type(input_value: Any) -> Tuple[Any, type]:
        return type(r := risky_cast(input_value)), r

    assert (int, 5) == val_and_type('5')
    assert (float, 4.55) == val_and_type('4.55')
    assert (str, 'foo') == val_and_type('foo')
    assert (dict, {'a': 5, 'b': 6}) == val_and_type("a=5, b=6")
    assert (bool, False) == val_and_type('False')
    assert (bool, True) == val_and_type('True')
    assert (dict, {'a': 5, 'b': True}) == val_and_type("a=5, b=True")


def test_extract_json():
    s: str = """The quick
    [[EMBEDDED_JSON_LINE]]{"color": "red", "value": "#f00"}
    Brown fox jumped 
    [[EMBEDDED_JSON_LINE]]{"color": "green", "value": "#0f0"}
    [[EMBEDDED_JSON_LINE]]{"color": "blue", "value": "#00f"}
    Over the foobar
    [[EMBEDDED_JSON_LINE]]{"color": "cyan", "value": "#0ff"}
    [[EMBEDDED_JSON_LINE]]{"color": "black", "value": "#000"}
    """
    extracted = parse_string_with_embedded_json(s)
    assert extracted.iloc[1, 0] == 'green'
    assert extracted.iloc[1, 1] == '#0f0'


def foobar(x: int) -> int:
    return x * 10


def starbar(x: int, y: int) -> str:
    return f"x={x} and y={y} so {x*y=}"


def test_process_queue():
    inputs: List[int] = list(range(25))
    star_inputs: List[Tuple[int, int]] = [(x, x + 2) for x in list(range(10))]

    result = process_queue(foobar, inputs, num_workers=5, cuts_per_worker=1, batching=False)
    assert result[:3] == [0, 10, 20]
    result = process_queue(foobar, inputs, num_workers=5, cuts_per_worker=1, batching=True)
    assert result[:3] == [0, 10, 20]
    result = process_queue(starbar, star_inputs, num_workers=5, cuts_per_worker=1, batching=False)
    assert result[:3] == ['x=0 and y=2 so x*y=0', 'x=1 and y=3 so x*y=3', 'x=2 and y=4 so x*y=8']
    try:
        print(process_queue(starbar, star_inputs, num_workers=5, cuts_per_worker=1, batching=True))
    except NotImplementedError as e:
        print(f"... not implemented yet")


def test_parsing():
    log: str = """
    log-2018-08-31-20-07-29:2018-08-31 20:07:24.748	[P2P2]	INFO protocol_handler.inl:1171	[1;33m[80.241.216.213:18080 OUT]  Synced 49156/1651224 (0.282567 sec, 353.898367 blocks/sec), 99.548294 MB [49255:o][0m
    log-2018-08-31-20-07-29:2018-08-31 20:07:25.040	[P2P2]	INFO protocol_handler.inl:1171	[1;33m[80.241.216.213:18080 OUT]  Synced 49256/1651224 (0.291497 sec, 343.056704 blocks/sec), 99.046715 MB [49355:][0m
    log-2018-08-31-20-07-29:2018-08-31 20:07:25.280	[P2P2]	INFO protocol_handler.inl:1171	[1;33m[80.241.216.213:18080 OUT]  Synced 49356/1651224 (0.239763 sec, 417.078532 blocks/sec), 98.918274 MB [49455:oo][0m
    log-2018-08-31-20-07-29:2018-08-31 20:07:25.644	[P2P2]	INFO protocol_handler.inl:1171	[1;33m[80.241.216.213:18080 OUT]  Synced 49456/1651224 (0.363321 sec, 275.238701 blocks/sec), 98.767830 MB [49555:o.][0m
    log-2018-08-31-20-07-29:2018-08-31 20:07:25.981	[P2P2]	INFO protocol_handler.inl:1171	[1;33m[80.241.216.213:18080 OUT]  Synced 49556/1651224 (0.335829 sec, 297.770592 blocks/sec), 98.502403 MB [49655:...][0m
    log-2018-08-31-20-07-29:2018-08-31 20:07:26.264	[P2P2]	INFO protocol_handler.inl:1171	[1;33m[80.241.216.213:18080 OUT]  Synced 49656/1651224 (0.281597 sec, 355.117420 blocks/sec), 98.384758 MB [49755:oo....][0m
    log-2018-08-31-20-07-29:2018-08-31 20:07:26.776	[P2P2]	INFO protocol_handler.inl:1171	[1;33m[80.241.216.213:18080 OUT]  Synced 49756/1651224 (0.404512 sec, 247.211455 blocks/sec), 98.054283 MB [49855:oo.....][0m
    log-2018-08-31-20-07-29:2018-08-31 20:07:27.066	[P2P2]	INFO protocol_handler.inl:1171	[1;33m[80.241.216.213:18080 OUT]  Synced 49856/1651224 (0.289512 sec, 345.408826 blocks/sec), 98.258347 MB [49955:oo.....][0m
    """

    # Set the rules for extracting the vectors
    record_delimiter: str = 'lo'
    tokens_dictionary: Dict[str, Tuple[str, str]] = {
        # varname: (left_token, right_token)
        'date_time_stamp': ('g-', ':'),
        'height': ('Synced ', '/'),
        'time_to_load': ('(', ' sec'),
    }

    # Extract the text into a timeseries
    timeseries: VectorMultiset = VectorMultiset(
        data=parse_string_with_manual_tokens(log,
                                             record_delimiter=record_delimiter,
                                             parallelize_processing=True,
                                             tokens_dictionary=tokens_dictionary))
    print(timeseries.data)


def mock_hill(kwargs) -> float:
    x = kwargs.get('x')
    y = kwargs.get('y')
    return 1 / x + 1 / (y ** 2)


def test_recursive_batch_evaluation():
    inital_point: Dict[str, Any] = {'x': 400, 'y': 300}
    current_best_input, history = recursive_batch_evaluation(
        func=mock_hill,
        initial_input=inital_point,
        selection_method=max,
        batch_generator_kwargs={'width_prct': 10, 'num_samples': 3},
        max_deep=5,
        print_progress=True,
        num_workers=64,
        return_history=True,
    )
    assert current_best_input['x'] < inital_point['x']
    assert current_best_input['y'] < inital_point['y']


def test_selector():
    inputs: List[Any] = ['abc', 'def', 'efg']
    eval_outputs: List[Dict[str, Any]] = [{'foobar': 'xyz', 'fitness': 4}, {'fitness': 99, 'baz': 12345},
                                          {'fitness': -5}]
    best = return_best_input(inputs, eval_outputs)
    assert best == 'def'


def y(x: float) -> float:
    return 10 - (x - 5) ** 2


def y_return_scalar(d) -> float:
    return y(d.get('x'))


def test_recursive_batch_evaluation_new_with_scalar(show_plot: bool = True):
    # Set up the run
    start_at: float = 15

    best_x, run_data = recursive_batch_evaluation(
        func=y_return_scalar,
        initial_input={'x': start_at},
        max_deep=20,
        return_history=True,
        evaluate_initial_inputs=True,
        infinite_memory=True,
        selection_method=max,
        catch_exceptions=False,
        batch_generator_kwargs={
            "width_prct": 50,
            "num_samples": 5,
            "fields": ['x'],
            "width_temperature_prct": 0,
        },
    )

    print(f"{best_x=}\n{run_data}")
    if show_plot:
        x_vals: List[float] = list(range(-10, 25))
        plt.figure(facecolor='w')
        plt.plot(x_vals, [y_return_scalar({'x': x}) for x in x_vals], color='k')
        x_vec = [d.get('in').get('x') for d in run_data]
        y_vec = [x.get('out') for x in run_data]
        plt.scatter(x_vec, y_vec, c=list(range(len(x_vec))), s=50)
        plt.show()

    assert best_x['x'] < start_at


def y_return_dict(d) -> dict:
    return {'y': y(d.get('x'))}


def test_recursive_batch_evaluation_new_with_dict(show_plot: bool = True):
    # Set up the run
    start_at: float = 15

    best_x, run_data = recursive_batch_evaluation(
        func=y_return_dict,
        initial_input={'x': start_at},
        max_deep=5,
        return_history=True,
        evaluate_initial_inputs=True,
        infinite_memory=True,
        selection_method='y',
        catch_exceptions=False,
        batch_generator_kwargs={
            "width_prct": 50,
            "num_samples": 5,
            "fields": ['x'],
            "width_temperature_prct": 0,
        },
    )

    print(f"{best_x=}\n{run_data}")
    if show_plot:
        x_vals: List[float] = list(range(-10, 25))
        plt.figure(facecolor='w')
        plt.plot(x_vals, [y_return_dict({'x': x})['y'] for x in x_vals], color='k')
        x_vec = [d.get('in').get('x') for d in run_data]
        y_vec = [x.get('out').get('y') for x in run_data]
        plt.scatter(x_vec, y_vec, c=list(range(len(x_vec))), s=50)
        plt.show()

    print(f"{run_data[-1]['in']}")
    assert run_data[-1]['in']['incremented_counter'] == 4
    assert run_data[-1]['in']['clock_time_at_start'] > 100000
    assert best_x['x'] < start_at
