#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys, getopt, os
from numpy import *
from matplotlib.pyplot import *
from copy import deepcopy
from scipy.stats.mstats import mquantiles

set_printoptions(precision=4,linewidth=500)

class Data:
    """ Class containing all data, names, values"""

    def __init__(self,X,Y, ShotNo, names, norm):
	"""construct a new Dir class with the specified directory"""
	self.X = double(X)
	self.Y = Y
	self.shots = ShotNo
	self.names = array(names)
	self.norm = array(norm)
	self.N = len(self.X)
	self.Ndim = shape(X)[1]
	self.remove_wrong()

    def __getitem__(self, vars):
	return self.X[:,self.ind(vars)].ravel()

    def remove_dim(self,ind):
	if  type(ind[0]) == str:
	    ind = in1d(self.names, [ind])
	if all(~ind):
	    return
	self.X = self.X[:,~ind]
	self.names = self.names[~ind]
	if len(self.norm) == self.Ndim :
	    self.norm = self.norm[~ind]
	self.Ndim = shape(self.X)[1]

    def ind(self,vars):
	ind =  where(in1d(self.names, array(vars)))[0]
	if len(ind) == 1:
	    return int(ind)
	else:
	    return ind

    def remove_data(self,ind):
	ind = ravel(ind)
	if len(ind) == 0:
	    return

	if dtype(ind[0]) == int:
	    ind = in1d(self.shots, ind)

	if sum(ind) < 2:
	    return

	self.X = self.X[~ind,:]
	self.Y = self.Y[~ind]
	self.shots = self.shots[~ind]
	self.N = len(self.X)


    def add_data(self,X,Y):
	self.X = vstack([self.X, X])
	self.Y = concatenate([self.Y, Y])
	self.shots = concatenate([self.shots, nan*ones(shape(Y))])
	self.N = len(self.X)

    def normalize(self, norm):
	""" renormalize data in all dimensions to be similar """
	self.norm = norm
	self.X -= self.norm[0,:]
	self.X /= self.norm[1,:]

	#for i in range(self.Ndim):
	    #self.X[self.X[:,i] > 0,i] /= self.norm[1,i]
	    #self.X[self.X[:,i] < 0,i] /= self.norm[2,i]

    def denormalize(self):
	""" renormalize data in all dimensions to be similar """
	self.X *= self.norm[1,:]
	self.X += self.norm[0,:]
	self.norm[0,:] = 0
	self.norm[[1,2],:] = 1


    def get_norm(self):
	norm = zeros((3,self.Ndim))
	for i in arange(self.Ndim):
	    d = self.X[:,i]
	    if median(d) != 0:
		m = median(d)
	    else:
		m = mean(d)
	    var = d-m;
	    sig1 =   mean( abs(var[var >= 0]))
	    sig2 =   mean( abs(var[var < 0]))
	    norm[:,i] = [m,sig1, sig2]
	return norm
	#self.norm = norm;

    def remove_wrong(self):
	"""  remove wrong shots (semi manual way ...) """
	wrong = list()
	wrong.append(array([ 3111, 3112, 3971, 1151, 1108, 2135, 2134, 2133 , 2699, 3046, 2544 , 3706, 4169, 4146, 3943, 3940, 3930,3905, 3768, 3767, 4201, 4300, 2138,2919, 4366, 4360, 1162, 3165  ] )) #something wrong
	wrong.append(arange(4374, 4398))  #missing diagnistics
	wrong.append(arange(4443, 4447)) #stabilozace
	wrong.append(concatenate([arange(3040, 3068), arange(1227, 1272), arange(1152, 1159)]))  #reverse field
	wrong.append(concatenate([arange(3800, 4020), arange(4476,4481)]))
	wrong = concatenate(wrong)
	self.remove_data(wrong)

    def __str__(self):
	print "========DATA OBJECT START======"
	print 'N', self.N, 'Ndim', self.Ndim
	print "size X", shape(self.X)
	print "size Y", shape(self.Y)
	print "size shots", shape(self.shots)
	print "size norm", shape(self.norm)
	print  "names", self.names
	print "mean", mean(self.X,0)
	print " max", amax(abs(self.X),0)
	return "========DATA OBJECT END======"