from MinMax import *
from math import log, exp, sin, cos, pi, sqrt
from simStorage import *

import Image, ImageDraw

def doTest():
    return SimToImg("213_hw\\poisson\\low_viscous_small_tstep")

class SimToImg:
    
    def __init__(self, strSubDir):
        if(strSubDir[-1] != '\\' and strSubDir[-1] != '/'):
            strSubDir = strSubDir+'/'
        self.psi = SimStorage(strSubDir)
        self.u = SimStorage(strSubDir)
        self.v = SimStorage(strSubDir)
        self.bScaleSet = False

    def numToStr(self, nNum):
        strOut = `nNum`
        while(len(strOut) < 3):
            strOut = '0'+strOut
        return strOut

    def grabFile(self, number):
        self.psi.convertData('psi',self.numToStr(number))
        self.u.convertData('u',self.numToStr(number))
        self.v.convertData('v',self.numToStr(number))
        self.psi.merge(self.u).merge(self.v)
        return self

    def makeImage(self):
        self.nX = self.psi.nX
        self.nY = self.psi.nY
        imOut = Image.new('RGB', (self.nX,self.nY))
        ranges = self.psi.myMinMax.getRange()
        scale = ranges[3]*ranges[3]+ranges[4]*ranges[4]
        scale = sqrt(scale)
        scale = 30.0/scale
        if(not self.bScaleSet):
            self.bScale = scale
        scale = self.bScale
        draw = ImageDraw.Draw(imOut)
        def MakeSharpLineFunc(fRange, nVals):
            return lambda x: 4*int(16-16*sqrt((sin(2*pi*nVals*x/fRange))**2))
        def MakePsiFunc(minMax, nCnt):
            fMin = minMax.m_min[2]
            fMax = minMax.m_max[2]
            fRange = fMax-fMin
            return MakeSharpLineFunc(fRange,nCnt)
        psiFunc = MakePsiFunc(self.psi.myMinMax, 10)
        psiArr = self.psi.formattedTriples
        # print `self.nX`+","+`self.nY`
        # print imOut.size
        for i in range(0, self.nX):
            for j in range(0, self.nY):
                draw.point((i,j), fill=(255, 255-psiFunc(psiArr[j][i][2]), 255-psiFunc(psiArr[j][i][2])))
        for i in range(11,self.nX, 11):
            for j in range(11, self.nY, 11):
                u = int(scale*self.u.formattedTriples[j][i][2])
                v = int(scale*self.v.formattedTriples[j][i][2])
                draw.ellipse([i-1,j-1,i+1,j+1],fill=(0,0,255))
                draw.line([(i,j), (i+u,j+v)], fill=(0,0,255))
        del draw
        return imOut