import matplotlib.pyplot as plt
import numpy as np
import random
import copy


def get_spike_sequence(size=1000, loc=12.5, spike_range_lower=50, spike_range_upper=200,
                       spike_location_lower=300, spike_location_upper=1000, if_need_location=False):
    if spike_location_upper >= size:
        spike_range_upper = size - 1
    if spike_location_lower >= spike_location_upper:
        spike_location_upper = spike_location_upper * 0.4
    location = random.randint(spike_location_lower, spike_location_upper)
    seq = np.random.normal(loc=loc, scale=1.5, size=size)
    seq[location] = random.randint(spike_range_lower, spike_range_upper)
    if if_need_location:
        return seq, location
    else:
        return seq


def get_dual_spike_sequence(size=1000, loc=12.5, downward_spike=2, upward_spike=2, midpoint=500, if_need_location=False):
    if midpoint > 0.8 * size:
        midpoint = int(size / 2)
    seq = np.random.normal(loc=loc, scale=1.5, size=size)
    downward_loc = random.randint(int(0.1*size), midpoint)
    upward_loc = random.randint(midpoint, int(0.9 * size), )

    seq[downward_loc] = seq[downward_loc] - loc * downward_spike * random.randint(9, 15) / 10
    seq[upward_loc] = seq[upward_loc] - loc * upward_spike * random.randint(9, 15) / 10
    if if_need_location:
        return seq, [downward_loc, upward_loc]
    else:
        return seq


def get_step_shift_sequence(size=1000, loc_list=None):
    if not loc_list:
        loc_list = [10, 50, 30, 40]
    block_length = int(size / len(loc_list))
    result = []
    for loc in loc_list:
        result += list(np.random.normal(loc=loc, scale=1.5, size=block_length))
    return np.array(result)


def get_slow_slop_sequence(size=1000, initial_value=10, is_upward=True):
    loc_list = [copy.deepcopy(initial_value) for _i in range(10)]
    for i in range(40):
        if is_upward:
            initial_value += random.randint(-12, 25) / 10
        else:
            initial_value -= random.randint(-12, 25) / 10
        loc_list.append(initial_value)
    return get_step_shift_sequence(size, loc_list)


def get_periodic_sequence(size=1000, seed_list=None):
    if not seed_list:
        seed_list = [0, 5, 8, 9, 8, 5, 0, -5, -8, -9, -8, -5, 0] * 5
    for i in range(len(seed_list)):
        seed_list[i] += random.randint(-20, 20) / 10
    return get_step_shift_sequence(size, seed_list)


def gen_random_sequence(size):
    result_dict = {
        'spike_up': get_spike_sequence(size=size, loc=12.5, spike_range_lower=50, spike_range_upper=200,
                                       spike_location_lower=300, spike_location_upper=1000),
        'spike_down': get_spike_sequence(size=size, loc=12.5, spike_range_lower=-200, spike_range_upper=-50,
                                         spike_location_lower=300, spike_location_upper=1000),
        'step_up': get_step_shift_sequence(size=size, loc_list=[10, 50]),
        'step_down': get_step_shift_sequence(size=size, loc_list=[100, 50]),
        'slop_up': get_slow_slop_sequence(size=size, initial_value=10, is_upward=True),
        'slop_down': get_slow_slop_sequence(size=size, initial_value=100, is_upward=False),
    }
    seed = random.randint(0, len(result_dict) - 1)
    return result_dict, seed


def gen_full_sequence(size, is_random=False):
    result_dict = {
        'spike': get_dual_spike_sequence(size=size, loc=12.5, downward_spike=4, upward_spike=4, midpoint=int(size/2)),
        'step_up': get_step_shift_sequence(size=size, loc_list=[10, 50]),
        'step_down': get_step_shift_sequence(size=size, loc_list=[100, 50]),
        'slop_up': get_slow_slop_sequence(size=size, initial_value=10, is_upward=True),
        'slop_down': get_slow_slop_sequence(size=size, initial_value=100, is_upward=False),
    }
    temp_list = list(result_dict.values())
    if is_random:
        random.shuffle(temp_list)
    seq_list = []
    for item in temp_list:
        seq_list += list(item)
    return np.array(seq_list)


