import pickle

import urllib2

import matplotlib
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import matplotlib.cbook as cbook
from mpl_toolkits.mplot3d import Axes3D

from volume_data import SparseVolumeData
from volume_data import FullVolumeData
from volume_data import get_entrez_ids

def trim(sparse_volume_data, keep=0.001):
    """Trim a SparseVolumeData object and return the trimmed SparseVolumeData object."""
    n = sparse_volume_data.size()
    keep_size = int(n * keep)
    values = sorted(sparse_volume_data.value)
    keep_above = values[n - keep_size]
    print 'min:', min(sparse_volume_data.value), ', max:', max(sparse_volume_data.value), ', keep_size:', keep_size, ', keep_above:', keep_above
    keep_values = values[(n - keep_size):]
    #print 'keep_values:', keep_values
    ''' '''
    trimmed_data = SparseVolumeData(sparse_volume_data.entrez_id, sparse_volume_data.symbol)
    trimmed_data.x = []
    trimmed_data.y = []
    trimmed_data.z = []
    trimmed_data.value = []
    for i in range(n): 
        if sparse_volume_data.value[i] >= keep_above:
            trimmed_data.x.append(sparse_volume_data.x[i])
            trimmed_data.y.append(sparse_volume_data.y[i])
            trimmed_data.z.append(sparse_volume_data.z[i])
            trimmed_data.value.append(sparse_volume_data.value[i])
    return trimmed_data
    
def get_hit_total(symbol):
    url = 'http://neuinfo.org:8080/services/rest/federation/search?query=' + symbol
    f = urllib2.urlopen(url)
    lead100 = f.read(100)
    #print 'symbol: ' + symbol + ', response: ' + lead100
    hit_total = lead100.split('total="')[1].split('"')[0]
    return hit_total

def get_trim_percent():
    trim_percent = 75.0 # default

    # read list of gene symbols and entrez gene ids from text file
    f = open('trimpercent.txt', 'r')
    i = 0
    while True:
        line = f.readline()
        if line == '':
            break
        if line.startswith('#'):
            continue
        line = line.strip()
        trim_percent = float(line)
        #print 'trim percent: ' + str(trim_percent)
        break
    f.close()
    return trim_percent
        
def scatter_3d(x, y, z, values, xmax, ymax, zmax, title='title'):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    scat = ax.scatter(x, y, z, s=values, c=values, marker='o')

    ax.text2D(0.05, 0.95, title, transform=ax.transAxes)
    
    ax.set_xlim3d(0, xmax)
    ax.set_ylim3d(0, ymax)
    ax.set_zlim3d(0, zmax)

    ax.set_xlabel('X Coordinate')
    ax.set_ylabel('Y Coordinate')
    ax.set_zlabel('Z Coordinate')

    fig.colorbar(scat, shrink=0.5, aspect=5)

    plt.show()

'''
start main script
''' 

# get list of entrez ids     
entrez_ids = get_entrez_ids()
print 'entrez_ids:', entrez_ids

# get trim percent, convert to decimal
trim_value = get_trim_percent() / 100

# initialize max coordinates across all genes
overall_max_x = 0
overall_max_y = 0
overall_max_z = 0

# create list for trimmed sparse volume objects
trimmed_sv_list = []

entrez_ids_str = ''
symbols_str = ''

# for each gene (entrez_id), get sparse volume object
for entrez_id in entrez_ids:

    # unpickle object 
    filename = 'sparseVolume' + str(entrez_id) + '.pkl'
    print 'Unpickling ' + filename + ' ...'
    pkl_file = open(filename, 'rb')
    sparse_volume_data = pickle.load(pkl_file) 
    pkl_file.close()
    #print 'Printing sparse volume object to string ...'
    #print 'sparse volume data:\n', sparse_volume_data.to_string()
    print 'sparse volume info:: ' + sparse_volume_data.to_string_info()
    
    # build up string of entrez_ids
    entrez_ids_str += sparse_volume_data.entrez_id + ' '
    symbols_str += sparse_volume_data.symbol + ' '
    
    # trim
    trimmed_sv_data = trim(sparse_volume_data, trim_value)
    #print 'trimmed sparse volume info:: ' + trimmed_sv_data.to_string_info()
    
    # get max coords
    max_x = trimmed_sv_data.max_x()
    if max_x > overall_max_x:
        overall_max_x = max_x
    max_y = trimmed_sv_data.max_y()
    if max_y > overall_max_y:
        overall_max_y = max_y
    max_z = trimmed_sv_data.max_z()
    if max_z > overall_max_z:
        overall_max_z = max_z
    print 'max x,y,z: ' + str(max_x) + ',' + str(max_y) + ',' + str(max_z)
    
    # get hit total
    hit_total = get_hit_total(sparse_volume_data.symbol)
    print 'hit total: ' + str(hit_total)
    
    # scale hit total for plot
    hit_total = float(hit_total) / 1000
    print 'scaled hit total: ' + str(hit_total)
    
    # replace expression values with single hit total
    trimmed_sv_data.value = [hit_total] * trimmed_sv_data.size()
    #print 'trimmed_sv_data.value: ' + str(trimmed_sv_data.value)
    
    # add this trimmed sparse volume object to list
    trimmed_sv_list.append(trimmed_sv_data)

    print '\n'
    
print 'overall max x,y,z: ' + str(overall_max_x) + ',' + str(overall_max_y) + ',' + str(overall_max_z)

# create 3d array based on overall max x,y,z and initialized to integer zero value
full_volume_data = FullVolumeData()
full_volume_data.xmax = overall_max_x
full_volume_data.ymax = overall_max_y
full_volume_data.zmax = overall_max_z
full_volume_data.data = [[[0 for z in range(overall_max_z + 1)] for y in range(overall_max_y + 1)] for x in range(overall_max_x + 1)]

# for each trimmed sparse volume object (now with hits data)
for trimmed_sv_hits in trimmed_sv_list:
    #print 'trimmed_sv_hits: ' + trimmed_sv_hits.to_string_detail() + '\n'
    
    # aggregate hits at each 3d coordinate
    for i in range(trimmed_sv_hits.size()):
        full_volume_data.data[trimmed_sv_hits.x[i]][trimmed_sv_hits.y[i]][trimmed_sv_hits.z[i]] += int(trimmed_sv_hits.value[i])

#print 'full_volume_data: ' + full_volume_data.to_string()

# convert 3D array into sparse volume object
x_coords = []
y_coords = []
z_coords = []
values = []
for x in range(full_volume_data.xmax + 1):
    for y in range(full_volume_data.ymax + 1):
        for z in range(full_volume_data.zmax + 1):
            if full_volume_data.data[x][y][z] > 0:
                x_coords.append(x)
                y_coords.append(y)
                z_coords.append(z)
                values.append(full_volume_data.data[x][y][z])
plot_sv_data = SparseVolumeData(entrez_ids_str, symbols_str)
plot_sv_data.x = x_coords
plot_sv_data.y = y_coords
plot_sv_data.z = z_coords
plot_sv_data.value = values
print 'plot_sv_data info:: ' + plot_sv_data.to_string_info()
#print 'plot_sv_data detail:: ' + plot_sv_data.to_string_detail()

# round to nearest tens ceiling
xmax = int((float(plot_sv_data.max_x()) / 10) + 1) * 10
ymax = int((float(plot_sv_data.max_y()) / 10) + 1) * 10
zmax = int((float(plot_sv_data.max_z()) / 10) + 1) * 10
print 'x,y,z max: ' + str(xmax) + ',' + str(ymax) + ',' + str(zmax)

# plot
title = 'Genes: ' + symbols_str
scatter_3d(plot_sv_data.x, plot_sv_data.y, plot_sv_data.z, plot_sv_data.value, xmax, ymax, zmax, title)

