import random
import time
import copy

import numpy as np
from skimage.restoration import estimate_sigma

from .expr import *
from .fractdim import *


class IFS3(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=1)
                      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)]

    @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):
        return ";".join([str(e) for e in self._expr])

    def __copy__(self):
        ret = IFS3(nexpr=self._nexpr, world=self._world)
        ret._expr = [copy.copy(e) for e 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(self._expr).mutation()

    def step(self):
        r,g,b,a = [int(e)
                   for e in self._world[self._position[0]][self._position[1]]]
        arg = r
        arg <<= 8
        arg += g
        arg <<= 8
        arg += b
        arg <<= 8
        arg += a
        arg <<=16
        arg += self._position[0]
        arg <<=16
        arg += self._position[1]

        ret = int(random.choice(self._expr).eval(arg))
        dr, dg, db, da = r,g,b,a

        #self._position[1] = (ret & 0xFFFF)%self.width
        self._position[1] = ((ret & 0xFFFF)*self.width)//0x10000
        ret >>= 16
        self._position[0] = ((ret & 0xFFFF)*self.height)//0x10000
        ret >>= 16
        a = ret & 0xFF
        ret >>= 8
        b = ret & 0xFF
        ret >>= 8
        g = ret & 0xFF
        ret >>= 8
        r = ret & 0xFF

        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
        score = sum(scores)/len(scores)


        score += mod 
        if sigma and sigma > 0:
            score -= sigma/100
        colscore = abs(colcount-1024) / 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