candle_generator.py

import pygame
import time
import random

wants_to_quit = False
pause = False
delay = 1.0
plus_pressed = False
minus_pressed = False
screen_size = (640, 360)
screen = pygame.display.set_mode(screen_size, pygame.RESIZABLE)

BLACK = (0, 0, 0)
GREY = (127, 127, 127)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
DARK_RED = (127, 0, 0)
GREEN = (0, 255, 0)
DARK_GREEN = (0, 127, 0)

FPS = 60.0


def parse_timestamp(data):
    result = {}
    for key in data:
        frame_arr = key.split(":")
        frame_num = int(frame_arr[0]) * 30 + int(frame_arr[1])
        result[frame_num] = data[key]
    return result


timestamp = parse_timestamp({
    "0:00": "sideways",
    "0:26": "dump",
    "4:02": "sideways",
    "4:22": "dump",
    "6:20": "pump",
    "11:01": "major dump",
    "14:11": "pump",
    "16:20": "sideways",
    "18:20": "dump",
    "20:20": "pump"
})


def get_current_state(frame_num):
    last = "sideways"
    for num in timestamp:
        if num > frame_num:
            return last
        else:
            last = timestamp[num]
    return last


def generate_candle(o, avg=None, state="sideways"):
    inc = 0
    if "sideways" in state:
        inc = 0

    if random.randint(0,10) > 6:
        if "dump" in state:
            inc = random.random()*10
        elif "pump" in state:
            inc = -random.random()*10

    if "major" in state:
        inc *= 2

    c = o + (random.random() * 2 - 1) - inc * 0.5
    h = max(o, c) + pow(random.random(), 2)*2 + abs(inc)*random.random()
    l = min(o, c) - pow(random.random(), 2)*2 - abs(inc)*random.random()
    return o, h, l, c


def map_value(val, min_val=-5, max_val=5):
    return (1 - (val - min_val) / (max_val - min_val)) * screen_size[1]


def draw_candle(index, o, high, low, c, min_val=-5, max_val=5):
    y = map_value(max(o, c), min_val, max_val)
    w = 8
    h = abs(map_value(o, min_val, max_val) - map_value(c, min_val, max_val))
    x = index * (w + 2)
    rect = (x, y, w, h)
    x2 = x + w * 0.5

    COLOR = GREEN if o < c else RED
    COLOR2 = DARK_GREEN if o < c else DARK_RED

    pygame.draw.aaline(screen, COLOR, (x2, map_value(high, min_val, max_val)), (x2, map_value(low, min_val, max_val)))
    screen.fill(COLOR2, rect)
    pygame.draw.rect(screen, COLOR, rect, 1)


def get_data_from(prices):
    all = []
    for o in range(4):
        all += [item[o] for item in prices]

    min_val = min(all)
    max_val = max(all)
    avg_val = (min_val + max_val) * 0.5
    return min_val, max_val, avg_val


prices = []
o = 0
avg = 0
for i in range(int(screen_size[0] / 10)):
    price = generate_candle(o, avg)
    prices.append(price)

    min_val, max_val, avg_val = get_data_from(prices)
    avg = avg_val
    o = price[3]

start_time = time.time()
frame_num = 0
while not wants_to_quit:
    elapsed = abs(time.time() - start_time)
    start_time = time.time()

    state = get_current_state(frame_num)
    screen.fill(BLACK)

    min_val, max_val, avg_val = get_data_from(prices)

    for i in range(len(prices)):
        o, h, l, c = prices[i]
        draw_candle(i, o, h, l, c, min_val, max_val)

    if not pause:
        prices = prices[1:]
        prices.append(generate_candle(prices[-1][3], avg_val, state))
        print(state)

    pygame.display.flip()

    events = pygame.event.get()
    for event in events:
        if event.type == pygame.KEYUP or event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                wants_to_quit = True
            elif event.key == pygame.K_SPACE and event.type == pygame.KEYUP:
                pause = not pause
        elif event.type == pygame.VIDEORESIZE:
            screen_size = event.size
            screen = pygame.display.set_mode(screen_size, pygame.RESIZABLE)
    time.sleep(max(0, 1.0 / FPS - elapsed))

    pygame.image.save(screen, "images/frame_%d.png" % frame_num)

    frame_num += 1

    if frame_num >= 900:
        break
pygame.quit()