import random
import time
import copy

import numpy as np
from skimage.restoration import estimate_sigma

from .expr import *
from .fractdim import *


class IFS5(object):
    """ @brief 1 Variable IFS with R,G,B,A X,Y decomposition """

    #height=width=1024
    height=width=768
    #height=width=512

    def __init__(self, nexpr=4, init_sz=1, world=None):
        self._nexpr = nexpr
        self._expr = [[RpnExpr(sz=init_sz, nvar=6) for _ in range(6)]
                      for _ in range(nexpr)]
        self._world = world
        if self._world is None:
            self._world = self.get_world()
        self._position = [random.randint(0, self.height-1),
                          random.randint(0, self.width-1)]
        self._col = [random.randint(0, 255) for _ in range(4)]

    @classmethod
    def get_world(cls):
        return np.zeros((cls.height,cls.width,4), dtype=np.uint8)

    def raz_world(self):
        for i in range(self.height):
            for j in range(self.width):
                self._world[i][j] = [0,0,0,0]
    
    def __str__(self):
        el_lbl = 'XYrgba'
        fel = lambda el: ';'.join(['%s=(%s)' % (el_lbl[i], str(el[i]))
                                   for i in range(len(el))])
        ret = ';'.join(['<%s>' % fel(el) for el in self._expr])
        return ret

    def __copy__(self):
        ret = IFS5(nexpr=self._nexpr, world=self._world)
        ret._expr = [[copy.copy(e) for e in el]
                      for el in self._expr]
        #ret._world = copy.deepcopy(self._world)
        return ret

    def get_image(self):
    	return self._world
 
    def mutation(self, n=1, rand=True):
        n = 1 if n <= 1 else random.randint(1,n)
        for _ in range(n):
            random.choice(random.choice(self._expr)).mutation()

    def step(self):
        dr,dg,db,da = [int(e)
                       for e in self._world[self._position[0]][self._position[1]]]

        cx = int((self._position[0] / self.width) * ((1<<64)-1))
        cy = int((self._position[0] / self.height) * ((1<<64)-1))

        args = (cx,cy,dr,dg,db,da)
        rx,ry,r,g,b,a = (int(expr.eval(*args))
                         for expr in random.choice(self._expr))
        rx = (rx * (self.width-1)) // ((1<<64)-1)
        ry = (ry * (self.height-1)) // ((1<<64)-1)
        self._position[0], self._position[1] = rx, ry
        r,g,b,a = (c%256 for c in (r,g,b,a))

        sa = a/255
        outa = (a + da*(1-sa))/255
        if outa == 0:
            ro,go,bo = 0,0,0
        else:
            ro, go, bo = [(c*(a/255)+dc*da*(1-sa))/outa
                          for c, dc in ((r,dr), (g, dg), (b, db))]
        r,g,b,a = [int(e) for e in (ro,go,bo,outa*255)]


        self._world[self._position[0]][self._position[1]] = [r,g,b,a]

    def score(self):
        start = time.time()
        
        colcount = len(np.unique(np.reshape(self._world[:,:,:3], (-1,3)),
                       axis=0))

        #sigma = estimate_sigma(self._world, multichannel=True, average_sigmas=True)
        sigmas = estimate_sigma(self._world[:,:,:3], multichannel=True,
                                average_sigmas=False)

        scores = [fractal_dimension(self._world[:,:,i]*self._world[:,:,3]/255)
                  for i in range(3)]
        #alpha score
        #scores += [fractal_dimension(self._world[:,:,3])]

        gray = rgb2gray(self._world)
        graysigma = estimate_sigma(gray)
        grayscore = fractal_dimension(gray)
        del(gray)

        sigmas += [graysigma]*3
        sigmas = [0 if np.isnan(sigma) else sigma for sigma in sigmas]

        scores += [grayscore]*3

        sigma = sum(sigmas)/len(sigmas)

        mod = abs(scores[0]-scores[1])
        mod += abs(scores[0]-scores[2])
        mod += abs(scores[0]-scores[3])
        mod += abs(scores[1]-scores[2])
        mod += abs(scores[1]-scores[3])
        mod /= 5

        null_comp = 0
        for i in range(3):
            null_comp += 1 if scores[i] == 0 else 0

        if null_comp >= 2:
            score = 0
            mod *= 0.8
        else:
            score = sum(scores)/len(scores)

        score += mod 

        if sigma and sigma > 0:
            score -= sigma/100
        #colscore = abs(colcount-1024) / 1e5
        colscore = abs(colcount-2048) / 1e5
        score -= colscore

        printscore = lambda arr: '['+(', '.join(['%1.3f' % e for e in arr]))+']'
        print("colscore %3.3f (%4d colors) scores time %5.2fs" % (colscore,
                                                                 colcount,
                                                                 time.time() - start))
        print("SIGMAS : %s SIGMA  : %f " %  (printscore(sigmas), sigma))
        print("SCORES : %s SCORE : %r" % (printscore(scores), score))
        return score