#!/usr/bin/python
from numpy import *
import numpy as np
import scipy.stats as st
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import copy
import os.path
import time
import re
import sys

if(len(sys.argv) != 3):
   print "Usage:"
   print "paretoPlot.py <output_postfix> <input_dir>"
   exit(1)

dirName = "run"
dirName2 = "mmm"
paretoFileName = "pareto.dat"
outputDir = sys.argv[1]
inputDir = sys.argv[2]
colorNr = 0
step = 250


def MatrixCreate( width, heigth ):
   "Creates a matrix with these dimensions"
   matrix = zeros((width, heigth))
   return matrix

def VectorCreate(width):
   "Creates a vector"
   return zeros((width), dtype='f')

def getDirs(directory):
   dirs = []
   for name in os.listdir(directory):
      runDir = directory + "/" + name
      result = re.match('.*' + dirName + '.*', name)
      if os.path.isdir(runDir) & (result is not None):
         for dir2 in os.listdir(runDir):
            result = re.match('.*' + dirName2 + '.*', dir2)
            if os.path.isdir(runDir + "/" + dir2) & (result is not None):
               dirs.append(runDir + "/" + dir2)
   return dirs

def countLines(filename):
   modularityFile = open(filename, 'r')
   lines = 0
   for line in modularityFile:
      lines+=1
   modularityFile.close()
   return lines

def countNumbers(filename):
   numbers = 0
   modularityFile = open(filename, 'r')
   line = modularityFile.readline()
   for number in line.split():
      numbers+=1
   modularityFile.close()
   return numbers

def readModularityFile(filename, statistics, fileNr):
   modularityFile = open(filename, 'r')

   currentLine = 0
   currentNumber = 0

   modularityFile.seek(0)
   for line in modularityFile:
      line = line.strip()
      for number in line.split():
         statistics[fileNr, currentLine, currentNumber] = float(number)
         currentNumber+=1
      currentNumber = 0
      currentLine+=1

   modularityFile.close()

def printMatrix(matrix):
   print("Printing matrix:")
   for row in matrix:
      for value in row:
         print str(value)

def printMatrix(statistics):
   print("Printing matrix:")
   for matrix in statistics:
      for row in matrix:
         print row

def writeToFiles(statistics):
   shape = statistics.shape
   lines = shape[1]
   numbers = shape[2]
   for number in toPlot:
      f = open(outputDir + "/" + fileNames[number]+".dat", 'w')
      for lineNr in range(lines):
         f.write(str(statistics[:,lineNr,number].mean()) + " ")
         f.write(str(statistics[:,lineNr,number].var()) + " ")
         f.write(str(st.cmedian(statistics[:,lineNr,number])) + " ")
         f.write(str(statistics[:,lineNr,number].min()) + " ")
         f.write(str(statistics[:,lineNr,number].max()) + " ")
         f.write('%6.3f %6.4f ' % st.normaltest(statistics[:,lineNr,number]))
         f.write('%6.4f '% st.sem(statistics[:,lineNr,number]))
         f.write("\n")
      f.close()

def plotPoint(firstValue, secondValue, values):
   plt.scatter(firstValue, secondValue, s=(values+1)*8, c=0.1, cmap=cm.Spectral)

def plotParetoFront(dataMatrix, fileNr, plotNr):
   print "Writing figure: " + "run_" + str(fileNr)+ "/gen_" + str((plotNr-1)) +".png"
   plt.xlabel(r'Fitness in % correct', fontsize = 16)
   plt.ylabel(r'Connection-cost in sum-squared length', fontsize = 16)
   plt.ylim((-10,1000))
   plt.xlim((0,1))
   plt.axhline(100, color='g', linestyle='--')
   plt.scatter(dataMatrix[0], dataMatrix[1], s=dataMatrix[3], c=dataMatrix[2], cmap=cm.Spectral, vmin = 1, vmax = 1000, alpha=0.5)
   cbar=plt.colorbar(ticks=[1, 1000])
   cbar.ax.set_yticklabels([1, 1000])
   cbar.set_label(r'Number of Individuals')
   plt.savefig(outputDir + "/" + "run_" + str(fileNr) + "/gen_" + str((plotNr-1)) +".png")
   plt.clf()

def readParetoFile(filename, fileNr):
   modularityFile = open(filename, 'r')

   plot = 0
   plotNr = 0

   previousFirstValue = 0
   previousSecondValue = 0
   values = 1
   previousValue=False

   dataMatrix = [[],[],[], []]
   for line in modularityFile:
      line = line.strip()
      nr = 0
      for number in line.split():
         if(nr==0):
            plotNr=int(number)
         elif(nr==2):
            firstValue=float(number)
         elif(nr==3):
            secondValue=abs(float(number))
         nr+=1
      if(plot!=plotNr):
         if((plotNr-1)%step == 0):
            plotParetoFront(dataMatrix, fileNr, plotNr)
            dataMatrix = [[],[],[], []]
            values = 1
            previousFirstValue = 0
            previousSecondValue = 0
            plot=plotNr
            previousValue=False

      if(plotNr%step == 0):
         if(previousValue):
            if((previousFirstValue == firstValue) & (previousSecondValue == secondValue)):
               values += 1
            else:
               dataMatrix[0].append(previousFirstValue)
               dataMatrix[1].append(previousSecondValue)
               dataMatrix[2].append(values)
               dataMatrix[3].append(values+5)
               values = 1
         else:
            previousValue=True
         previousFirstValue = firstValue
         previousSecondValue = secondValue
   plotParetoFront(dataMatrix, fileNr, plotNr)


def run2():
#   toPlot = [2, 3, 4, 19, 20, 21]
   if not os.path.exists(outputDir):
      os.makedirs(outputDir)

   dirs = getDirs(inputDir)
   files = len(dirs)

   fileNr = 0
   for name in dirs:
      print("Reading " + name+"/" + paretoFileName + "...")
      if not os.path.exists((outputDir + "/" + "run_" + str(fileNr))):
         os.makedirs((outputDir + "/" + "run_" + str(fileNr)))
      readParetoFile(name+"/" + paretoFileName, fileNr)
      fileNr+=1

run2()
