#!/usr/bin/env python3
# Copyright 2019 Yann Weber <yannweb@member.fsf.org>
#
# pyasciimandel is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# any later version.
#
# pyasciimandel is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with pyasciimandel.  If not, see <http://www.gnu.org/licenses/>.
#
import sys
import os
import random
import curses
import multiprocessing


def in_set(x0, y0, colors, max_iter):
    """ Return True if given point is in mandel set"""
    c = y0 * 1j + x0
    z = 0.0j+0.0
    i = 0
    for i in range(max_iter):
        z = z**2 + c
        if z.imag ** 2 + z.real ** 2 >= 4:
            return colors[int(i * len(colors) / max_iter)]
    return ' '


def process_line(args):
    point_list, colors, max_iter = args
    return ''.join([in_set(x, y, colors, max_iter) for y, x in point_list])


def fill_screen_multi(stdscr, pool, zoom, center, colors, max_iter):

    lines, cols = stdscr.getmaxyx()
    lines-=1
    term_sz = lines, cols
    set_len = (2, 3)

    deltas = [ set_len[i] / zoom / term_sz[i] for i in range(2)]
    ranges = [ (int(-sz/2), int(sz/2)) for sz in term_sz]

    ys, xs = [ [center[c] + (i * deltas[c]) for i in range(*ranges[c])]
               for c in range(2)]

    points = [([(y, x) for x in xs], colors, max_iter) for y in ys]


    for i, line in enumerate(pool.imap(process_line, points)):
        stdscr.addstr(i,0, line)


if __name__ == '__main__':

    center = [0,-0.5] # Mandelbrot Set center (y,x)
    # other color sets
    #std_colors = ".,:;-=+*oO8&@#"
    #std_colors = ".:-=+*#%@"
    #std_colors = '.\'`^",:;-=+*mwqpdbkhaoOIlXYUJCLQ0OZ><?][}{)(|\\/#MW&8%B@$'
    std_colors = '.\'`^",:;-=+*mwaoO><?][}{)(|\\/#MW&8%B@$'
    colors = list(std_colors)

    max_iter = 290
    zoom = 1
    decimal_mode = False

    pool = multiprocessing.Pool(os.cpu_count())
    stdscr = curses.initscr()
    curses.noecho()
    curses.cbreak()
    try:
        while True:
            iter2cols = lambda i, max_i: colors[int(i * len(colors) / max_iter)]
            fill_screen_multi(stdscr, pool, zoom, center, colors, max_iter)
            stdscr.refresh()
            k = stdscr.getkey()
            if k == '\x1b':
                if stdscr.getkey() != '[':
                    continue
                stdscr.refresh()
                k = stdscr.getkey()
                kmap = {'D':'j', 'A':'i', 'B':'k', 'C':'l'}
                if k in kmap:
                    k = kmap[k]
                else:
                    #print('\\1b[%r' % k, file=sys.stderr)
                    #sys.stderr.flush()
                    continue

            delta_xmove = 0.7 / zoom
            delta_ymove = 0.5 / zoom
            delta_zoom = 1.2
            
            if k == 'q':
                exit(0)
            elif k == '+':
                zoom *= delta_zoom
            elif k == '-':
                if zoom >= delta_zoom:
                    zoom /= delta_zoom
            elif k == 'j':
                center[1] -= delta_xmove
            elif k == 'l':
                center[1] += delta_xmove
            elif k == 'i':
                center[0] -= delta_ymove
            elif k == 'k':
                center[0] += delta_ymove
            elif k == 's':
                random.shuffle(colors)
            elif k == 'S':
                colors = list(std_colors)
            elif k == 'd':
                decimal_mode = not decimal_mode
            elif k == 'p':
                max_iter += 25
            elif k == 'm':
                max_iter -= 25
            elif k == 'I':
                lines, cols = stdscr.getmaxyx()
                win = curses.newwin(20, 40, (lines - 20)//2,
                                    (cols - 40)//2)
                win.addstr(1,12, 'Current status')
                win.addstr(2,12, '==============')
                win.addstr(5,0, '''\
 Max iteration = %d
 Zoom level    = %f

 X = %f
 Y = %f''' % (max_iter, zoom, center[0], center[1]))
                win.border()
                win.refresh()
                k = stdscr.getkey()
            else:
                lines, cols = stdscr.getmaxyx()
                win = curses.newwin(22, 40, (lines - 20)//2,
                                    (cols - 40)//2)
                win.addstr(1,11, 'pyasciimandel')
                win.addstr(2,11, '=============')
                win.addstr(4,0, '''\
 +   zoom in
 -   zoom out
 j   left
 l   right
 i   up
 k   down
 s   shuffle colors
 S   standard colors
 d   decimal mode (useless)
 p   add 25 to maximum iterations
 m   sub 25 to maximum iterations
 I   information
 q   quit
 h   display this help
''')
                win.addstr(20, 1, 'Press any key to continue')
                win.border()
                win.refresh()
                k = stdscr.getkey()
                if k == 'q':
                    exit(0)

    finally:
        curses.nocbreak()
        curses.echo()
        curses.endwin()
    exit()