import io
import random

from PIL import Image
from flask import Flask, request, jsonify
import matplotlib.image as mpimg
import numpy as np
import numba
import pickle
import base64

app = Flask(__name__)


@numba.njit
def optimized_mandelbrot(n_rows, n_columns, iterations, cx, cy):
    x_cor = np.linspace(-2, 2, n_rows)
    y_cor = np.linspace(-2, 2, n_columns)
    output = np.zeros((n_rows,n_columns))
    c = cx + 1j * cy
    for i in range(n_rows):
        for j in range(n_columns):
            z = x_cor[i] + y_cor[j] * 1j
            count = 0
            for k in range(iterations):
                z = (z*z) + c
                count += 1
                if np.abs(z) > 4:
                    break
            output[i, j] = count
    return output.T


def open_image_as_array(path):
    return mpimg.imread(path)


def get_encoded_img(arr):
    img = Image.fromarray(arr).convert("L")
    img_byte_arr = io.BytesIO()
    img.save(img_byte_arr, format='PNG')
    encoded_img = base64.encodebytes(img_byte_arr.getvalue()).decode('ascii')
    return encoded_img


@app.route('/getImage', methods=['GET'])
def get_image():
    with open('data_100.pickle', 'rb') as f:
        images = pickle.load(f)
    buffered = io.BytesIO()
    random.choice(images).save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue())
    return jsonify({
        "code": 0,
        "image": str(img_str.decode('ascii')),
        "first_time": 1
    })


def start_app():
    app.run()


if __name__ == '__main__':
    start_app()