import matplotlib
# Force matplotlib to not use any Xwindows backend.
matplotlib.use('Agg')
from matplotlib import cm
import math as m
import numpy as np
import scipy
import pylab as pl
import os
import sys
import time
import pwd
import glob
from sklearn import neural_network as neunet                         

#print(matplotlib.pyplot.colormaps())
#quit()

deg = np.pi/180.0

fbas = ""
mode = ""
opts = ""

if len(sys.argv)>=3:
    fbas = sys.argv[1]
    mode = sys.argv[2]
    if len(sys.argv)>=4:
        opts = sys.argv[3]
else:
    print(' usage: '+sys.argv[0]+' FILE_BASE MODE [ OPTS ]\n')
    print(' FILE_BASE text data \n')
    print(' MODE can be tricky. \n')
    print(' OPTS  plotquery ... (see code) \n')
    quit()

fids = []
fnoids = []
if True:
    for fil in glob.glob(fbas+'*objlis.txt'):
        if '+id_' in fil: 
            fids.append(fil)
        else: 
            fnoids.append(fil)

print('# files w/id, no id:',len(fids),len(fnoids))

ns = len(fids)+len(fnoids)
nx = 20
ny = 20

data = np.zeros((ns,nx*ny))
labels = np.zeros(ns)
names = np.zeros(ns,dtype=object)

def lablabs(s):
#    if s=='gax': return 4
#    elif s=='glcl': return 2 
#    elif s=='opcl': return 1 
    if s=='gax': return 1
    elif s=='glcl': return 1 
    elif s=='opcl': return 1 
    return 0;


gaxtypes = 'Galaxy LINER Radio LSB_G EmG GiC GinPair StarburstG AGN'
starcltypes = 'GlCl OpCl'

def datatoimg(brthis,gmagthis,x,y):
    H, xedges, yedges = np.histogram2d(brthis,gmagthis,bins=(x, y))
    H = H.T; H = H.flatten()
    return H,xedges,yedges

def gentrain(BRBR,GG,nt,x,y):
    # this generates a set of random draws from gmag, br = GG, BRBR
    # also, shifting gmag as if at a different distance
    nx = len(x)-1; ny = len(y)-1
    dtarg = np.zeros((nt,nx*ny))
    for i in range(nt):
        dgmag = 1*np.random.uniform()-0.5
        gmagthis = np.copy(GG)+dgmag
        msk = gmagthis<21.0
        gmagthis = gmagthis[msk]
        brthis = BRBR[msk]
        if len(brthis)>20:
            frac = ((0.9*np.random.uniform()+0.1)*len(brthis))
            msk = np.random.uniform(0,1,len(brthis))<frac
            brthis = brthis[msk]; gmagthis = gmagthis[msk]
        H, xedges, yedges = datatoimg(brthis,gmagthis,x,y)
        dtarg[i] = H
    return dtarg

def getinfo(fil,x,y):
    with open(fil) as myfile:
        head = [next(myfile) for x in range(20)]
    name,otype = '',''; label = 0.0
    for lin in head:
        if (' id: ' in lin) and (not 'OBJNAME' in lin):
            name,otype = lin.split(':')[1:3]
            name = name.strip(); otype = otype.strip(); 
            if "NAME " in name: name = name[5:]
            if otype=='Cl*' and 'NGC 4214' in name:
                name = 'NGC 4214'; otype = 'Galaxy'
            elif '2XMM J140312.4+542054' in name:
                otype = 'Galaxy'
            elif 'NGC  2366' in name:
                name = 'NGC 2366'
                otype = 'Galaxy'
            if otype in gaxtypes:
                #print(name,otype)
                label = lablabs('gax')
                break
            elif otype in starcltypes:
                #print(name,otype)                
                label = lablabs(otype.lower())
                break
            name = ''; otype=''
            # print(name,otype)
    cols=(6,7)
    gmag,br = np.loadtxt(fil,unpack=True,usecols=cols,max_rows=10000)
    H, xedges, yedges = datatoimg(br,gmag,x,y)
    return label, name, otype, H, xedges, yedges, gmag, br

brblue, brred = -1.0, 2.0
gbright, gfaint = 14.0, 22.0
x = np.linspace(brblue,brred,nx+1)
y = np.linspace(gbright,gfaint,ny+1)

ctr = 0
trainingobj = ["Bootes", "Coma Dwarf", "CVn I dSph"]
print('training on this list:',trainingobj)
ftrainingobj= []

for f in fids:
    labels[ctr],names[ctr],otype,data[ctr],xe,ye,gmag,br = getinfo(f,x,y)
    for targ in trainingobj:
        if targ in names[ctr]:
            ftrainingobj.append(f)
            break
    ctr += 1
nids=ctr

for f in fnoids:
    labels[ctr],names[ctr],otype,data[ctr],xe,ye,gmag,br = getinfo(f,x,y)
    ctr += 1

# data = data.reshape((ns,nx*ny))

# sort :
#idx = np.argsort(labels)[::-1]
#labels = labels[idx]; data = data[idx]

ntbg = 1000
bgdtrain = np.zeros((ntbg,nx*ny))
bgltrain = np.zeros(ntbg)
# create training data
bgdata = np.mean(data[labels==0,:],axis=0)
for i in range(ntbg):
    for j in range(nx*ny):
        bgdtrain[i,j] = np.random.poisson(bgdata[j])
#ntbg = ctr-nids
#bgdtrain = data[nids:]
#bgltrain = labels[nids:]

ntsig = 900
ntuse = int(ntsig/len(ftrainingobj))
ntsig = int(ntuse*len(ftrainingobj))
sigdtrain = np.zeros((0,nx*ny))
sigltrain = np.zeros((0,nx*ny))
for f in ftrainingobj:
    label,name,otype, H, xe, ye, gmag, br = getinfo(f,x,y)
    HH = gentrain(br,gmag,ntuse,x,y)
    LL = np.ones(ntuse)*label
    sigdtrain = np.append(sigdtrain,HH,axis=0)
    sigltrain = np.append(sigltrain,LL)

# dtrain = np.append(bgdtrain[:],sigdtrain[:],axis=0)
# ltrain = np.append(bgltrain,sigltrain,axis=0)

dtrain = np.append(sigdtrain[:],bgdtrain[:],axis=0)
ltrain = np.append(sigltrain,bgltrain)

for i,d in enumerate(data):
    d /= np.amax(d)
    data[i] = d

for i,d in enumerate(dtrain):
    d /= np.amax(d)
    dtrain[i] = d