import pygame
from pygame.locals import *
from math import *

def mandelPlot(c):
    x = c.real
    y = c.imag
    cx = x
    cy = y
    n = 0
    while x*x + y*y <= 2 and n < ITER_LIM:
        x,y = x*x - y*y + cx,2*x*y + cy
        n += 1
    return n

pygame.init()

WIN_W = 1280
WIN_H = 720
REDUC = 8
FPS = 60
ITER_LIM = 100

clock = pygame.time.Clock()
c_cold = (0,0,64)
c_med = (192,192,192)
c_hot = (255,192,0)
c_lime = (128,255,0)

disp = pygame.display.set_mode((WIN_W,WIN_H),0,32)
#frac = pygame.Surface((WIN_W,WIN_H))

pygame.display.set_caption("Mandelbrot set")
font = pygame.font.Font("courier.ttf",12)

#Control
(x,y) = (3/2,0)
zoom = 2 / min(WIN_W,WIN_H)
grab = False
(mx,my) = (-16,-16)
change = True
drawHD = False

#Main loop
done = False
while not done:
    if change:
        w = WIN_W if drawHD else WIN_W//REDUC
        h = WIN_H if drawHD else WIN_H//REDUC
        frac = pygame.Surface((w,h))
        frac.fill(c_cold)
        for j in range(h):
            for i in range(w):
                c = complex(x + (i - w//2) * zoom,y + (j - h//2) * zoom)
                plotColor = mandelPlot(c)
                if plotColor == ITER_LIM:
                    color = (0,0,0)
                else:
                    t = plotColor / ITER_LIM
                    color = (255*t,192*t,64*(1 - t))
                frac.set_at((i,j),color)

    change = False
    if drawHD:
        disp.blit(frac,(0,0))
        drawHD = False
    else:
        disp.blit(pygame.transform.scale(frac,(WIN_W,WIN_H)),(0,0))

    #UI
    text = font.render("Scale: " + str(zoom) + " units/px",True,c_lime)
    disp.blit(text,(8,8))
    if not grab:
        coord = font.render("c = " + str(x + (mx - WIN_W//2) * zoom) + " + " + str(y + (my - WIN_H//2) * zoom) + "i",True,c_lime)
        disp.blit(coord,(mx + 16,my - 16))

    #User interaction
    for event in pygame.event.get():
        if event.type == QUIT:
            done = True
        elif event.type == MOUSEMOTION:
            (ox,oy) = (mx,my)
            (mx,my) = event.pos
            if grab:
                x += (ox - mx) * zoom
                y += (oy - my) * zoom
                change = True
        elif event.type == MOUSEBUTTONDOWN:
            if event.button == 4:
                zoom *= 0.9
                change = True
            elif event.button == 5:
                zoom *= 1.2
                change = True
            elif event.button == 1 and not grab: grab = True
            elif event.button == 2 and not grab:
                x = x + (mx - WIN_W//2) * zoom
                y = y + (my - WIN_H//2) * zoom
                change = True
        elif event.type == MOUSEBUTTONUP and event.button == 1:
            grab = False
        elif event.type == KEYDOWN:
            if event.key == K_ESCAPE: done = True
            elif event.key == K_SPACE:
                drawHD = True
                change = True

    pygame.display.update()
    clock.tick(FPS)

#Done
pygame.quit()

const ColorRamp::Point & ColorRamp::find(float x, int & ind) const
{
    ind = 0;
    for(uint i=0; i<pts.size(); ++i)
    {
        if( pts[i].x >= x )
        {
            ind = i-1;
            break;
        }
    }
    return pts[ind];
}

const ColorRamp::Point & ColorRamp::findNext(int ind) const
{
    ind++;
    if(ind >= (int)pts.size())
        return pts.back();
    else
        return pts[ind];
}

Color ColorRamp::color(float x) const
{
    if(outside(x))
    {
        const Point & pt = outsidePoint(x);
        return pt.r | pt.g<<8 | pt.b<<16;
    }
   
    int ind = -1;
    const Point & p1 = find(x, ind);
    const Point & p2 = findNext(ind);
   
    float t = (x - p1.x) / (p2.x - p1.x);
    byte r = p1.r + (p2.r - p1.r) * t;
    byte g = p1.g + (p2.g - p1.g) * t;
    byte b = p1.b + (p2.b - p1.b) * t;
    return Color(r, g, b);

}