#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import args
import numpy as np
import matplotlib.pyplot as plt
from IPython import embed


class MultichannelStimulus(object):

    def __init__(self, namespace):
        assert(namespace is not None)
        assert(isinstance(namespace, argparse.Namespace))
        self._constrast_stimulus = None
        self._zap_stimulus = None
        self._phase_stimulus = None
        self._eod_ampl = namespace.amplitude
        self._eod_freq = namespace.frequency
        self._df = namespace.beatfrequency
        self._dcf = namespace.dcfrequency
        self._outfile = namespace.outfile
        self._dt = namespace.stepsize
        self._maxfreq = namespace.maxfreq
        self._zap_duration = 10.0
        self._contrast_sweep_duration = 5.0
        self._phase_sweep_duration = 5.0
        self.create_contrast_stim()
        self.create_phase_stim()
        self.create_zap_stim()
        

    def create_contrast_stim(self):
        t = np.arange(0, self._contrast_sweep_duration, self._dt)
        sweep_up = np.linspace(0., 1.0, len(t))
        am = np.sin((self._eod_freq + self._df) * 2 * np.pi * t) * 0.5
        dc = np.sin(2*np.pi*self._dcf*t) * 0.5
        am_sweep = am * sweep_up
        dc_sweep = dc * sweep_up
       
        self._contrast_stimulus = np.hstack((am_sweep, dc_sweep, (am + dc_sweep), (dc + am_sweep))) 
        header = {" ": "Contrast sweeps for ampullary and tuberous pathways",
                  "sd": 1.0,
                  "deltat": "%.5f s" % self._dt,
                  "eodf": self._eod_freq,
                  "df": self._df,
                  "dcf": self._dcf,
                  "T": "%.3f s" % (4 * self._contrast_sweep_duration)}
        self.save_stim(np.arange(len(self._contrast_stimulus)) * self._dt, self._contrast_stimulus, header,
                       "contrast_sweep")

    def create_phase_stim(self):
        t = np.arange(0, self._phase_sweep_duration, self._dt)
        phase_sweep = np.linspace(0., 2*np.pi, len(t))
        am = np.sin((self._eod_freq + self._df) * 2 * np.pi * t)
        dc = np.sin(2*np.pi*self._df*t + phase_sweep)

        self._phase_stimulus = (am + dc)/2.
        header = {" ": "Phase sweep stimulating ampullary and tuberous pathways",
                  "sd": 1.0,
                  "deltat": "%.5f s" % self._dt,
                  "eodf": self._eod_freq,
                  "df": self._df,
                  "T": "%.3f s" % (self._phase_sweep_duration)}
        self.save_stim(t, self._phase_stimulus, header, "phase_sweep")
    
    def create_zap_stim(self):
        t = np.arange(0, self._zap_duration, self._dt)
        m = self._maxfreq/self._zap_duration

        dc = np.sin(2*np.pi*m*t*t) * 0.5
        am = np.sin(2*np.pi*(m*t + self._eod_freq)*t) * 0.5
        self._zap_stimulus = np.hstack((dc, am, (dc + am)))
        header = {" ": "Zap stimulus for ampullary and tuberous pathways",
                  "sd": 1.0,
                  "deltat": "%.5f s" % self._dt,
                  "eodf": self._eod_freq,
                  "maxf": self._maxfreq,
                  "T": "%.3f s" % (3*self._zap_duration)}
        self.save_stim(np.arange(len(self._zap_stimulus))*self._dt, self._zap_stimulus, header, "zap")

    def save_stim(self, time, stimulus, header, suffix):
        assert(isinstance(header, dict))

        with open('_'.join((self._outfile, suffix)) + '.dat', "w") as f:
            for h in header.items():
                if len(h[0].strip()) == 0:
                    f.write("# %s\n" % (h[1]))
                else:
                    f.write("# %s = %s\n" % (h[0], h[1]))
            f.write("\n")
            f.write("#Key\n")
            f.write("# t \t x \n")
            f.write("# s \t 1 \n")

            for t, s in zip(time, stimulus):
                f.write("%.5f\t %.7f\n" % (t, s))
            f.write("\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=args.description)
    parser.add_argument('-f', '--frequency', metavar="f", type=float, default='100',
                        help='The fish\'s EOD frequency')
    parser.add_argument('-a', '--amplitude', metavar='a', type=float,
                        help='The fish\'s EOD amplitude in mV/cm')
    parser.add_argument('-bf', '--beatfrequency', default=20, metavar='b', type=float,
                        help=args.beat_help)
    parser.add_argument('-dc', '--dcfrequency', default=20, metavar='dcfreq', type=float,
                        help=args.dc_help)
    parser.add_argument('-dt', '--stepsize', type=float, default=1./20000,
                        help=args.dt_help)
    parser.add_argument('-mf', '--maxfreq', type=float, default=100.,
                        help=args.maxfreq_help)
    parser.add_argument('outfile', help="prefix name of the output file may include the path")

    args = parser.parse_args()

    stim = MultichannelStimulus(args)