#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
from .lase_simu import LaseSimu


class SpectrumSimu(LaseSimu):
    
    def __init__(self): 
        n = 4096
        super(SpectrumSimu, self).__init__(n)
        
        self.waveform_size = n
        self.lase_base = LaseSimu(self.waveform_size)
                
        self.spectrum = np.zeros(self.lase_base.sampling.n, dtype=np.float32)
        self.demod = np.zeros((2,self.lase_base.sampling.n,n))
        
        self.demod[0,:] = 0.49 * (1-np.cos(2*np.pi*np.arange(self.lase_base.sampling.n)/self.lase_base.sampling.n))
        self.demod[1,:] = 0
        
        self.velocity_signal = 1
        
        self.reset()
        
    def set_velocity(self, velocity, SNR):
        """ Generates the velocity beat signal for velocimeter simulation
        
            Args:
                - velocity: Object velocity (m/s)
                - SNR: Signal-to-noise
        """
        time = np.linspace(0, self.model.sampling.n*self.model.sampling.dt, self.model.sampling.n)
        omega_doppler = 2*np.pi*velocity/self.model._wavelength # Doppler shift (rad/s)
        self.velocity_signal = np.cos(omega_doppler*time) + np.random.randn(self.model.sampling.n)/SNR
    
    def get_spectrum(self):
        self.update()
        self.set_velocity(velocity=4.2, SNR=0.25)
        
        adc = self.model.adc_from_voltage(
                self.model.photodetection_voltage(self.model._optical_attenuation 
                                                  * self._live_laser_power 
                                                  * self.velocity_signal),
                n_avg=1)
        self.spectrum = np.abs(np.fft.fft(adc))
        
    def set_demod(self):
        pass

@@ -123,20 +123,16 @@
 eSimu):%0A
-    
 %0A    def
@@ -147,17 +147,16 @@
 _(self):
- 
 %0A       
@@ -211,25 +211,17 @@
 it__(n)%0A
-        
 %0A
+
         
@@ -301,25 +301,9 @@
 ze)%0A
-                
 %0A
+
     
@@ -409,16 +409,17 @@
 eros((2,
+ 
 self.las
@@ -436,28 +436,21 @@
 pling.n,
+ 
 n))%0A
-        
 %0A       
@@ -463,16 +463,17 @@
 demod%5B0,
+ 
 :%5D = 0.4
@@ -491,52 +491,146 @@
 os(2
-*
+ * 
 np.pi
-*np.arange(self.lase_base.sampling.n)/
+ *%0A                                            np.arange(self.lase_base.sampling.n) /%0A                                            
 self
@@ -678,23 +678,16 @@
 d%5B1,
+ 
 :%5D = 0%0A
-        
 %0A   
@@ -716,25 +716,17 @@
 nal = 1%0A
-        
 %0A
+
         
@@ -738,25 +738,17 @@
 reset()%0A
-        
 %0A
+
     def 
@@ -856,25 +856,17 @@
 ulation%0A
-        
 %0A
+
         
@@ -1009,54 +1009,137 @@
 e(0,
- self.model.sampling.n*self.model.sampling.dt,
+%0A                           self.model.sampling.n *%0A                           self.model.sampling.dt,%0A                          
  sel
@@ -1187,23 +1187,52 @@
  = 2
-*
+ * 
 np.pi
-*
+ *%5C%0A                        
 velocity
 /sel
@@ -1227,17 +1227,19 @@
 velocity
-/
+ / 
 self.mod
@@ -1252,16 +1252,40 @@
 velength
+%0A                       
  # Doppl
@@ -1360,16 +1360,48 @@
 *time) +
+%5C%0A                              
  np.rand
@@ -1435,17 +1435,15 @@
 g.n)
-/
+ / 
 SNR%0A
-    
 %0A   
@@ -1539,25 +1539,17 @@
 R=0.25)%0A
-        
 %0A
+
         
@@ -1665,16 +1665,17 @@
 nuation 
+*
 %0A       
@@ -1712,26 +1712,24 @@
             
- *
  self._live_
@@ -1740,16 +1740,17 @@
 r_power 
+*
 %0A       
@@ -1791,18 +1791,16 @@
         
- *
  self.ve
@@ -1888,24 +1888,16 @@
 t(adc))%0A
-        
 %0A    def

clean spectrum_simu.py