#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# vim: expandtab shiftwidth=4
'''
Проект "Построение графика"
Графический редактор в возможностью рисовать графики
см. диаграммы PyGameState_*.dia
'''

from math import *
import sys

import pygame
pygame.init()
pygame.font.init()
# Векторы, указывающие на соседнюю точку
dxy=[(-1,-1),(-1,0),(-1,1),(0,1),(1,1),(1,0),(1,-1),(0,-1)]
dxyW=len(dxy)

def distance(M1,M2):
    'Расстояние от точки до точки'
    return sqrt(float(M1[0]-M2[0])**2+float(M1[1]-M2[1])**2)

def dxydist(M1, M2, M, R, ind):
    '''Вспомогательная функция. Вычисляет координаты точек, соседних с M,
    соответствующих векторам, индексы которых в dxy перечислены в ind;
    вычисляет, насколько отличается заданное значение R 
    от суммы расстояний от каждой из этих точек до центров эллипса.
    Возвращает список вида [ (отличие, точка, индекс в dxy), ...]'''
    return [(abs((R-distance(m,M1)-distance(m,M2))),m,i) for m,i in [((M[0]+dxy[i][0],M[1]+dxy[i][1]),i) for i in ind]]

def Ellipse(screen, M1, M2, M0, pen, width):
    'Рисование эллипса по двум центрам и точке M0'
    L=distance(M1,M2)
    R=distance(M0,M1)+distance(M0,M2)
    if int(R)<=int(L):
        D=[M1,M2]
    elif L<3:
        pygame.draw.circle(screen, pen, M1, max(R/2, width), width)
        return
    else:
        D=[M0]
        r,M,v=min(dxydist(M1, M2, M0, R, range(dxyW)))
        # здесь иногда вечный цикл: какая-то неточность в расчётах
        c=0
        while M!=M0 and c <10020:
            c+=1
            if c > 10000: print r,M,v, R, M1, M2
            D.append(M)
            r,M,v=min(dxydist(M1, M2, M, R, ((dxyW+v-1)%dxyW,v,(dxyW+v+1)%dxyW)))
    pygame.draw.polygon(screen, pen, D, width)

def scale(x, X0, X1, Z0, Z1):
    'Преобразует координату x в диапазоне [X0, X1] в координату z в диапазоне [Z0, Z1]'
    return Z0 + (x!=X0 and (Z1-Z0)*(x-X0)/float(X1-X0)) or 0

def Function(screen, M1, M2, pen, width):
    'Рисование графика функции fml в диапазоне fml(X0)..fml(X1)'
    global X0, X1, fml
    mw,mh=min(M1[0],M2[0]),min(M1[1],M2[1])
    Mw,Mh=max(M1[0],M2[0]),max(M1[1],M2[1])
    count=max(Mw-mw,2)
    dots=[(x,eval(fml)) for x in [scale(i,0,count-1,X0,X1) for i in range(count)]]
    Y0,Y1=min([e[1] for e in dots]), max([e[1] for e in dots])
    gr=[(round(scale(x,X0,X1,mw,Mw)),round(scale(y,Y0,Y1,mh,Mh))) for x,y in dots]
    pygame.draw.lines(screen, pen, False, gr, width)

def Draw(screen, pen, width, center, pos, figure, center2=None):
    if figure == "Line":
        pygame.draw.line(screen, pen, center, pos, width)
    elif figure == "Rect":
        pygame.draw.rect(screen, pen, (center, (pos[0]-center[0],pos[1]-center[1])), width)
    elif figure == "Function":
        Function(screen, center, pos, pen, width)
    elif figure == "Circle":
        r=int(round(distance(center, pos)))
        pygame.draw.circle(screen, pen, center, max(r, width), width)
    elif figure == "Ellipse":
        if center2:     # финальная фигура
            Ellipse(screen, center, center2, pos, pen, width)
        else:           # фигура второй стадии
            Ellipse(screen, center, (center[0]*.2+pos[0]*.8, center[1]*.2+pos[1]*.8), pos, pen, width)

def DrawState(screen, event, stages):
    'Рисование фигуры по точкам'
    global State, Stage, Center, Copy, pen, width, Action, Center2
    if event.type == pygame.MOUSEMOTION:
        if Stage == 2:
            screen.blit(Copy, (0,0))    # восстанавливаем 
            Draw(screen, pen, width, Center, event.pos, Action)
        if Stage == 3:
            screen.blit(Copy, (0,0))    # восстанавливаем 
            Draw(screen, pen, width, Center, event.pos, Action, Center2)
    elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
        if Stage == 1:                  # первое нажатие
            Stage += 1
            Center = event.pos
            Copy = screen.copy()        # состояние экрана без нарисованной линии
        elif Stage == stages:           # Последнее нажатие
            Stage = 1
        elif Stage == 2:                # второе нажатие из более чем двух
            Stage += 1
            Center2 = event.pos
    elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 3:
        if Stage > 1:                   # Нарисовано что-то резиновое
            screen.blit(Copy, (0,0))    # восстанавливаем 
            Stage = 1
    elif event.type == pygame.KEYDOWN and event.key == 27:
        screen.blit(Copy, (0,0))
        State = "Base"

def ScrInput(screen, prompt, pos):
    '''Ввести строку с клавиатуры, отображая её.
    Backspace -- удалить один символ, ESC -- вернуть None'''
    global Copy
    res = ""
    Copy = screen.copy()
    while True:
        text = font.render(prompt+": "+res, 1, (255,255,255))
        screen.blit(Copy, (0,0))
        screen.blit(text, pos)
        pygame.display.flip()
        event = pygame.event.wait()
        if event.type == pygame.KEYDOWN:
            if event.unicode:                     # real characters
                if event.unicode == u"\033":      # ESC
                    res = None
                    break
                elif event.unicode == "\010":     # Backspace
                    res=res[:-1]
                elif event.unicode in u"\r\n":    # Enter
                    break
                elif event.unicode > u" ":        # printable characters
                    res += event.unicode
                #print event.unicode.__repr__()
    screen.blit(Copy, (0,0))
    pygame.display.flip()
    return res

def Ask_n_Do(screen, what, prompt):
    '''Ввести строку, проверить, и сделать с ней что полагается.
    Вернуть True или False в зависимости от успешности'''
    global width, pen, fml, X0, X1
    str_res=ScrInput(screen,u"Введите "+prompt, (10,10))
    if str_res is None: return False    # ввод был прерван
    try:
        res = eval(str_res)
    except:     # неизвестно что, пусть это остаётся строкой
        res = str_res
    if what == "Width":
        if type(res) != int:
            print what, "must be integer"
            return False
        width = res
    elif what == "Color":
        if type(res) not in (list, tuple) or len(res) != 3:
            print what, "must be 3-element list"
            return False
        pen = [r%256 for r in res]      # лень проверять ещё и < 256
    elif what == "Absc":
        try:
            X0,X1=float(res[0]),float(res[1])
        except:
            print sys.exc_value
            return False
    elif what == "Funct":
        try:
            x=X0
            y=eval(str(res))
            fml=str(res)
            print "FML:",fml
        except:
            print sys.exc_value
            return False
    elif what == "Save":
        try:
            pygame.image.save(screen,res)
        except:
            print sys.exc_value
            return False
    elif what == "Edit":
        try:
            newscr = pygame.image.load(res)
            screen.blit(newscr,(0,0))
        except:
            print sys.exc_value
            return False
    else:
        print "Unknown", what
    return True

# Основной экран
W,H=640,480
screen=pygame.display.set_mode((W,H))

pen,width=(10,100,200),2
X0,X1=1,8
fml="sin(x)"
font = pygame.font.Font(None, W/24)
States=("Draw", "Input", "Base")
Figures={"o":"Circle", "l":"Line", "r":"Rect", "y":"Function"}
Figures3={"d":"Ellipse"}
Inputs={"w":    ("Width", u"ширину линии"),
        "c":    ("Color", u"цвет линии"),
        "f":    ("Funct", u"функцию для рисования"),
        "x":    ("Absc", u"минимальное, максимальное значение X"),
        "s":    ("Save", u"имя файла для сохранения"),
        "e":    ("Edit", u"имя файла для чтения"),
}
Center=(0,0)
Copy=None

State, Stage, Action = "Base", 1, "Line"
while True:
  pygame.display.flip()
  event = pygame.event.wait()
  # Выход по закрытию окна
  if event.type == pygame.QUIT: sys.exit()
  if State == "Base":   # основное состояние
    if event.type == pygame.KEYDOWN:
        key = event.unicode.lower()
        if key == "q":
            sys.exit()
        elif key in Figures.keys():     # Начать рисование фигуры
            State,Stage,Action = "Draw", 1, Figures[event.unicode.lower()]
            Copy = screen.copy()
        elif key in Figures3.keys():     # Начать рисование фигуры
            State,Stage,Action = "Draw3", 1, Figures3[event.unicode.lower()]
            Copy = screen.copy()
        elif key in Inputs.keys():      # Что-то ввести и обработать результат
            Ask_n_Do(screen, *Inputs[key])
  elif State == "Draw": # рисование фигуры
    DrawState(screen, event, 2)
  elif State == "Draw3": # рисование фигуры
    DrawState(screen, event, 3)