// $Id: pywrap.hoc,v 1.31 2012/08/04 03:19:13 samn Exp $ 

//* variables
declare("INITPYWRAP",0) // whether initialized properly

//* initialize pywrap
if(2!=name_declared("p")) {
  print "pywrap.hoc: loading python.hoc"
  load_file("python.hoc")
}
func initpywrap () { localobj pjnk
  INITPYWRAP=0
  if(2!=name_declared("p")){printf("initpywrap ERR0A: PythonObject p not found in python.hoc!\n") return 0}
  print p  
  pjnk=new PythonObject()
  if(!isojt(p,pjnk)){printf("initpywrap ERR0B: PythonObject p not found in python.hoc!\n")}
  if(!nrnpython("import numpy")) {printf("pypmtm ERR0C: could not import numpy python library!\n") return 0}
  INITPYWRAP=1
  return 1
}
initpywrap()

//** pypmtm(vec,samplingrate[,nw])
// this function calls python version of pmtm, runs multitaper power spectra, returns an nqs
obfunc pypmtm () { local sampr,spc,nw localobj vin,str,nqp,ptmp
  if(!INITPYWRAP) {printf("pypmtm ERR0A: python.hoc not initialized properly\n") return nil}
  if(!nrnpython("from mtspec import *")) {printf("pypmtm ERR0B: could not import mtspec python library!\n") return nil}  
/*  if(!nrnpython("import numpy")) {printf("pypmtm ERR0C: could not import numpy python library!\n") return nil}*/
  if(numarg()==0) {printf("pypmtm(vec,samplingrate)\n") return nil}
  vin=$o1 sampr=$2 str=new String()
  p.vjnk = vin.to_python()
  p.vjnk = p.numpy.array(p.vjnk)
  spc = 1.0 / sampr // "spacing"
  nw=4 if(numarg()>2) nw=$3
  sprint(str.s,"[Pxx,w]=mtspec(vjnk,%g,%d)",spc,nw)
  nrnpython(str.s)
  nqp=new NQS("f","pow")
  nqp.v.from_python(p.w)
  nqp.v[1].from_python(p.Pxx)
  return nqp
}

//** pybspow(vec,samplingrate[,maxf,pord])
// this function calls python version of bsmart, to get power pectrum, returns an nqs
// pord is order of polynomial -- higher == less smoothing. default is 12
obfunc pybspow () { local sampr,pord,maxf localobj vin,str,nqp,ptmp
  if(!INITPYWRAP) {printf("pybspow ERR0A: python.hoc not initialized properly\n") return nil}
  if(!nrnpython("from spectrum import ar")) {printf("pybspow ERR0B: could not import spectrum python library!\n") return nil}  
  if(numarg()==0) {printf("pybspow(vec,samplingrate)\n") return nil}
  vin=$o1 sampr=$2 str=new String()
  if(numarg()>2) maxf=$3 else maxf=sampr/2
  if(numarg()>3) pord=$4 else pord=64
  p.vjnk = vin.to_python()
  p.vjnk = p.numpy.array(p.vjnk)
  sprint(str.s,"Pxx,F=ar(vjnk,rate=%g,order=%d,maxfreq=%g)",sampr,pord,maxf)
  nrnpython(str.s)
  nqp=new NQS("f","pow")
  nqp.v[0].from_python(p.F)
  nqp.v[1].from_python(p.Pxx)
  return nqp
}

//** pyspecgram(vec,samplingrate[,orows])
// this function calls python version of specgram, returns an nqs
obfunc pyspecgram () { local sampr,spc,i,j,sz,f,tt,orows,a localobj vin,str,nqp,ptmp,vtmp
  if(!INITPYWRAP) {printf("pyspecgram ERR0A: python.hoc not initialized properly\n") return nil}
  if(!nrnpython("from matplotlib.mlab import specgram")) {printf("pyspecgram ERR0B: could not import specgram from matplotlib.mlab!\n") return nil}  
  if(numarg()==0) {printf("pyspecgram(vec,samplingrate)\n") return nil}
  a=allocvecs(vtmp)
  vin=$o1 sampr=$2 str=new String()
  if(numarg()>2)orows=$3 else orows=1
  p.vjnk = vin.to_python()
  p.vjnk = p.numpy.array(p.vjnk)
  sprint(str.s,"[Pxx,freqs,tt]=specgram(vjnk,Fs=%g)",sampr)
  nrnpython(str.s)
  if(orows) {
    {nqp=new NQS("f","pow") nqp.odec("pow")}
    {sz=p.Pxx.shape[0] nqp.clear(sz)}
    for i=0,sz-1 {
      {vtmp.resize(0) vtmp.from_python(p.Pxx[i]) f=p.freqs[i]}
      nqp.append(f,vtmp)
    }
  } else {
    nqp=new NQS("f","pow","t")
    sz = p.Pxx.shape[0]
    nqp.clear(sz * p.Pxx.shape[1])
    for i=0,sz-1 {
      {vtmp.resize(0) vtmp.from_python(p.Pxx[i]) f=p.freqs[i]}
      for j=0,vtmp.size-1 nqp.append(f,vtmp.x(j),p.tt[j])
    }
  }
  dealloc(a)
  return nqp
}

//** pycsd(vec1,vec2,samplingrate)
// this function calls python version of csd (cross-spectral density)
// returns an nqs with csd -- csd is non-directional
obfunc pycsd () { local sampr,a localobj v1,v2,str,nqp
  if(!INITPYWRAP) {printf("pycsd ERR0A: python.hoc not initialized properly\n") return nil}
  if(!nrnpython("from matplotlib.mlab import csd")) {printf("pycsd ERR0B: could not import csd from matplotlib.mlab!\n") return nil}  
  if(numarg()==0) {printf("pycsd(vec,samplingrate)\n") return nil}
  v1=$o1 v2=$o2 sampr=$3 str=new String()
  {p.vjnk1=v1.to_python() p.vjnk1=p.numpy.array(p.vjnk1)}
  {p.vjnk2=v2.to_python() p.vjnk2=p.numpy.array(p.vjnk2)}
  sprint(str.s,"[Pxy,freqs]=csd(vjnk1,vjnk2,Fs=%g)",sampr)
  nrnpython(str.s)
  nqp=new NQS("f","pow")
  nqp.v[0].from_python(p.freqs)
  nqp.v[1].from_python(p.Pxy)
  return nqp
}

//** pypsd(vec,samplingrate[,NFFT])
// this function calls python version of psd (power-spectral density)
// returns an nqs with psd
obfunc pypsd () { local sampr,NFFT localobj v1,str,nqp
  if(!INITPYWRAP) {printf("pypsd ERR0A: python.hoc not initialized properly\n") return nil}
  if(!nrnpython("from matplotlib.mlab import psd")) {printf("pypsd ERR0B: could not import psd from matplotlib.mlab!\n") return nil}  
  // nrnpython("from matplotlib.mlab import window_none")
  if(numarg()==0) {printf("pypsd(vec,samplingrate)\n") return nil}
  v1=$o1 sampr=$2 str=new String() 
  {p.vjnk1=v1.to_python() p.vjnk1=p.numpy.array(p.vjnk1)}
  if(numarg()>2) NFFT=$3 else NFFT=v1.size
  if(sz%2==1) sz+=1
  sprint(str.s,"[Pxx,freqs]=psd(vjnk1,Fs=%g,NFFT=%d)",sampr,NFFT)
  nrnpython(str.s)
  nqp=new NQS("f","pow")
  nqp.v[0].from_python(p.freqs)
  nqp.v[1].from_python(p.Pxx)
  return nqp
}

//** pycohere(vec1,vec2,samplingrate) 
// this function calls python version of cohere (coherence is normalized csd btwn vec1, vec2)
// returns an nqs with coherence
obfunc pycohere () { local sampr,a localobj v1,v2,str,nqp
  if(!INITPYWRAP) {printf("pycohere ERR0A: python.hoc not initialized properly\n") return nil}
  if(!nrnpython("from matplotlib.mlab import cohere")) {printf("pycohere ERR0B: could not import cohere from matplotlib.mlab!\n") return nil}  
  if(numarg()==0) {printf("pycohere(vec1,vec2,samplingrate)\n") return nil}
  v1=$o1 v2=$o2 sampr=$3 str=new String()
  {p.vjnk1=v1.to_python() p.vjnk1=p.numpy.array(p.vjnk1)}
  {p.vjnk2=v2.to_python() p.vjnk2=p.numpy.array(p.vjnk2)}
  sprint(str.s,"[Pxy,freqs]=cohere(vjnk1,vjnk2,Fs=%g)",sampr)
  nrnpython(str.s)
  nqp=new NQS("f","coh")
  nqp.v[0].from_python(p.freqs)
  nqp.v[1].from_python(p.Pxy)
  return nqp
}

//* pypca(matrix) - does PCA on input matrix and returns scores (projections onto PCs)
// rows of the matrix are observations, columns are 'features' or 'dimensions'
obfunc pypca () { local r,c,a localobj inm,inmT,vin,vout,str,mout
  if(!INITPYWRAP) {printf("pypca ERR0A: python.hoc not initialized properly\n") return nil}
  if(!nrnpython("from princomp import PCA")) {printf("pypca ERR0B: could not import PCA!\n") return nil}  
  if(!nrnpython("import numpy as np")){printf("pypca ERR0C: could not import numpy as np!\n") return nil}
  str=new String2()
  if(numarg()<1) {printf("pypca(Vector,rows,cols)\n") return nil}
  a=allocvecs(vin,vout)
  {inm=$o1 r=inm.nrow c=inm.ncol}
  inmT = inm.transpose() // transpose since to_vector goes in column ordering
  {vin.resize(r*c) inmT.to_vector(vin)}
  p.vjnk = vin.to_python() // convert to python format
  sprint(str.s,"vjnk=np.resize(vjnk,(%d,%d))",r,c)
  if(!nrnpython(str.s)){printf("pypca ERR0D: could not run %s\n",str.s) dealloc(a) return nil}
  if(!nrnpython("mypca=PCA(vjnk)")){printf("pypca ERR0E: could not run PCA\n") dealloc(a) return nil}
  if(!nrnpython("score=mypca.Y")){printf("pypca ERR0F: could not set scores\n") dealloc(a) return nil}
  sprint(str.s,"score=np.resize(score,(%d,1))",r*c)
  if(!nrnpython(str.s)){printf("pypca ERR0E: could not run %s\n",str.s) dealloc(a) return nil}
  vout.from_python(p.score) // convert to a hoc Vector  
  mout=new Matrix(c,r)//output as a matrix. NB: c,r are reversed from original for following transpose
  mout.from_vector(vout)//from_vector uses column ordering
  mout = mout.transpose()//so need to transpose
  dealloc(a)
  return mout
}

//* pyspecck(vec,sampr[,maxf,win]) - call ck's spectrogram.py
// vec = time-series. sampr = sampling rate (Hz).
// maxf = max frequency. win = window size (seconds) for specgram chunks
// system call to spectrogram.py file to display a spectrogram, writes temp
// file and then deletes it...
func pyspecck () { local i,sampr,maxf,win localobj fp,vec,str
  vec=$o1 sampr=$2
  if(numarg()>2)maxf=$3 else maxf=sampr/2
  if(numarg()>3)win=$4 else win=1
  str=new String2()
  fp=new File()
  if(!fp.mktemp()){printf("pyspecck ERR0: couldn't make temp file!\n") return 0}
  str.s=fp.getname()
  fp.wopen(str.s)
  for i=0,vec.size-1 fp.printf("%g\n",vec.x(i))
  fp.close()  
  sprint(str.t,"/usr/site/nrniv/local/python/spectrogram.py %s %g %g %g",str.s,sampr,maxf,win)
  print str.t
  system(str.t)
  fp.unlink()
  return 1
}

//* pykstest(vec1,vec2) - perform a two-sample, two-sided kolmogorov-smirnov test
// and return the p-value. kstest checks if values in vec1,vec2 come from same distribution (null hypothesis)
// returns -1 on failure. uses scipy.stats.ks_2samp function
func pykstest () { localobj v1,v2
  if(!INITPYWRAP) {printf("pykstest ERR0A: python.hoc not initialized properly\n") return -1}
  {v1=$o1 v2=$o2}
  if(!nrnpython("from scipy.stats import ks_2samp")) return -1
  {p.v1=v1.to_python() p.v2=v2.to_python()}
  if(!nrnpython("(D,pval)=ks_2samp(v1,v2)")) return -1
  return p.pval  
}