#    Grating Spectrometer software
#    Copyright (c) 2015 Joao Almeida  http://opentronix.com
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
    
import matplotlib.pyplot as plt
import numpy as np
import time;
import cv2
import os
CAP_PROP_BRIGHTNESS=10
CAP_PROP_CONSTRAST=11
CAP_PROP_HUE=13
CAP_PROP_EXPOSURE= 15
CAP_PROP_SATURATION=12
CAP_PROP_GAMMA=22

cap = cv2.VideoCapture(0)
cap.set(CAP_PROP_BRIGHTNESS,.5)
cap.set(CAP_PROP_CONSTRAST,.5)
cap.set(CAP_PROP_HUE,.5)
cap.set(CAP_PROP_SATURATION,0.2)

havezero=0

FPER="spectra"
OUTDIR="output"
try:
    os.stat(OUTDIR)
except:
    os.mkdir(OUTDIR)       

init=1
plt.ion()
plt.show()
while(True):
     # Capture frame-by-frame
    ret, frame = cap.read()

    
    key=cv2.waitKey(1) & 0xFF
    if key == ord('q'):
        break
    
    if (init):
        sy=frame.shape[0]
        yy = np.array( range(sy))
        yyc=.8611*yy+310.0492
        bgfreqr=yy*0
        bgfreqg=yy*0
        bgfreqb=yy*0
        bgfreqt=yy*0
    
    
    freqr=frame[:,:,2].sum(axis=1);freqr=freqr-min(freqr)
    freqg=frame[:,:,1].sum(axis=1);freqg=freqg-min(freqg)
    freqb=frame[:,:,0].sum(axis=1);freqb=freqb-min(freqb)
    freqt=freqr+freqg+freqb
    
    if key == ord('z'):
        bgfreqr=freqr
        bgfreqg=freqg
        bgfreqb=freqb
        bgfreqt=freqt
        print('zero')
        havezero=1
    if key == ord('x'):
        havezero=0
    if (havezero):
        freqr=freqr-bgfreqr+1000
        freqg=freqg-bgfreqg+1000
        freqb=freqb-bgfreqb+1000
        freqt=freqt-bgfreqt+1000
    
#    for i in yy:          
#        if freqr[i]>=freqg[i] and freqr[i]>=freqb[i]: freqt[i]=freqr[i]*3
#        if freqg[i]>=freqr[i] and freqg[i]>=freqb[i]: freqt[i]=freqg[i]*3
#        if freqb[i]>=freqg[i] and freqb[i]>=freqr[i]: freqt[i]=freqb[i]*3
    
    plt.clf()
        
    #plt.axis([0, yy[-1], 100, 4E5])
    plt.axis([yyc[0], yyc[-1], 100, 4E5])
    plt.yscale('log')

    plt.plot(yyc, freqt,color='gray')
    plt.plot(yyc, freqr,color='red')
    plt.plot(yyc, freqg,color='green')
    plt.plot(yyc, freqb,color='blue')
    
    
    #plt.axvline(180,color='blue');
    #plt.axvline(260,color='green');
    #plt.axvline(360,color='red');
    plt.axvline(390,color='purple');
    plt.axvline(470,color='blue');
    plt.axvline(525,color='green');
    plt.axvline(590,color='yellow');
    plt.axvline(624,color='red');
    plt.axvline(625,color='orange');
    
    
    
    plt.draw()
    
    frame[180,:,0]=255
    frame[260,:,1]=255
    frame[360,:,2]=255
    cv2.imshow('frame',frame)
    
    if key == ord('n'):
        FPER = raw_input("File name:")
    if key == ord('s'):
        fname=OUTDIR+"/"+FPER+"_"+str(time.time())
        plt.savefig (fname+".png",dpi=200)
        cv2.imwrite(fname+"_raw.bmp", frame)
        dt = np.dtype([('x', '|i'), ('l', 'd'), ('red', 'i'), ('green', 'i'),('blue', 'i'),('total', 'i'), ('bgred', 'i'), ('bggreen', 'i'),('bgblue', 'i'),('bgtotal', 'i')])
        a = np.zeros(len(yy), dt)
        a['x']=yy;a['l']=yyc;a['red']=freqr;a['green']=freqg;a['blue']=freqb;a['total']=freqt
        a['bgred']=bgfreqr;a['bggreen']=bgfreqg;a['bgblue']=bgfreqb;a['bgtotal']=bgfreqt
        np.savetxt(fname+'_data.txt', a, '%s',header="x l red green blue bgred bggreen bggblue",comments='')
        
        print "saved "+fname
    #time.sleep(.5)
    init=0
 # When everything done, release the capture
cap.release()
cv2.destroyAllWindows()