# -*- coding: utf-8 -*-
"""
Created on Wed Sep  5 09:52:56 2018

@author: lenovo
"""

import numpy as np
import math
import tensorflow as tf
import struct
import matplotlib
import matplotlib.pyplot as plt
from scipy import signal
import os

def LoadWave(filepath,offset,buf,buflen):
    f = open(filepath, 'rb')  
    ndx = 0
    f.seek(offset*2)
    while True:  
        temp=f.read(2)  
        if len(temp) == 2:  
            value = struct.unpack('h',temp)
            buf[ndx] = int(value[0])
            ndx += 1
            if (ndx == buflen):break
        else:  
            break  
    f.close()

def SaveWave(filepath,buf,repeat,append):
    if append:opt = 'ab' 
    else:opt = 'wb'

    f = open(filepath, opt)
    for n in range(repeat):
        buf= buf.astype(np.int16)
        f.write(buf)   
    f.close()

def SaveData(filepath,buf):
    f = open(filepath, 'wb')
    buf= buf.astype(np.float32)
    f.write(buf)   
    f.close()    
    
def LoadData(filepath,buf):    
    f = open(filepath, 'rb')  
    ndx = 0
    while True:  
        temp=f.read(4)  
        if len(temp) == 4:  
            value = struct.unpack('f',temp)
            buf[ndx] = float(value[0])
            ndx += 1
        else:  
            break    
    f.close()
    
def reset_graph(seed=42):
    tf.reset_default_graph()
    tf.set_random_seed(seed)
    np.random.seed(seed)
    
def show_one(wave,color):
    plt.title("wave", fontsize=14)
    instances = range(len(wave))
    plt.plot(instances, wave[instances], marker='.', markersize=1, c=color)
    plt.xlabel("test")   
    
#返回：每一行是x_delay个x对应1个y，总共batch_size行
def sample_get(wave_y,y_delay):
    ret_y = np.reshape(wave_y,[1,len(wave_y)])
    ret_yl = np.zeros([1,y_delay],dtype='float32')
    return ret_y,ret_yl

GlottisTime = 115
N = 30
M = 30  
batch_size = 1

def cond(n, x, yl,y,Wa,Wb):
    return n < GlottisTime

def body(n, x, yl,y,Wa,Wb):

    #x切片    
    x_delay = tf.reshape(tf.slice(x,[0,n],[batch_size,M]),[batch_size,M]) 
    
    #yl_切片
    y_delay = tf.reshape(yl,[batch_size,N])
        
    #递归方程计算出当前y值
    y_cur = tf.subtract(tf.matmul(x_delay,Wb),tf.matmul(y_delay,Wa)) 
    y_cur = tf.reshape(y_cur,[batch_size,1])
        
    #当前y值填充到输出序列中的第n个位置
    y_out = tf.concat([tf.slice(y,[0,0],[batch_size,n]),tf.to_float(y_cur)],axis=1) 
    y_out = tf.concat([y_out,tf.slice(y,[0,n+1],[batch_size,GlottisTime-n-1])],axis=1) 
    y_out.set_shape([batch_size,GlottisTime])
    
    #yl的0..N-1位向后移1位,当前输出值保存到yl的第0位
    yl = tf.concat([y_cur,tf.slice(yl,[0,0],[batch_size,N-1])],axis=1)    
    
    return n+1, x, yl,y_out,Wa,Wb
        
def main(param_mode,param_type,param_voice):
            
    #原始波形
    wave_y_base = np.zeros([GlottisTime],dtype='int16')  
    
    #加载y语音
    path = './data/std/{}{}.snd'.format(param_voice,GlottisTime)
    LoadWave(path,0,wave_y_base,GlottisTime)          
            
    reset_graph()
    ##### 输入层 #####

    Yd = tf.placeholder(tf.float32, [None,GlottisTime])       #目标输出
    Yl = tf.placeholder(tf.float64, [None,N])                 #Y历史输出
    Ki = tf.placeholder(tf.float32, [GlottisTime],name='Ki')  #FIR系数输入
    Wai = tf.placeholder(tf.float32, [N],name='Wai')  #IIR-a系数输入    
    Wbi = tf.placeholder(tf.float32, [M],name='Wbi')  #IIR-b系数输入    

    ###### FIR层 输入必须是32位#####
        
    if param_mode=='Train' and param_type=='FIR':
        K = tf.Variable(tf.zeros([GlottisTime]), dtype=tf.float32) 
        init_K_Ki = tf.assign(K,Ki)
    else:
        K = tf.cast(Ki,dtype=tf.float32)  
                            
    #因果滤波器0对称补零                        
    K_ = tf.concat([K,tf.zeros([GlottisTime-1])],axis=0)
    K_ = tf.reshape(K_,[2*GlottisTime-1,1,1]) 
    #构造脉冲 [ 1,GlottisTime-1个0 ] 
    P = tf.concat([tf.ones([1]),tf.zeros([GlottisTime-1])],axis=0)    
    P = tf.reshape(P,[1,GlottisTime,1])
    #卷积输出（模拟声门波）
    X = tf.nn.conv1d(P, K_, 1, 'SAME') 
    X = tf.reshape(X,[1,GlottisTime])
    X = tf.concat([tf.zeros([1,M-1]),X],axis=1)    #（前插M-1个0）    
    X = tf.cast(X,dtype=tf.float64) 
        
    ###### IIR层 输入必须是64位#####
    
    if param_mode=='Train' and param_type=='IIR' :
        Wbv = tf.Variable(tf.zeros([M]),dtype=tf.float32) 
        Wav = tf.Variable(tf.zeros([N]),dtype=tf.float32)
        init_W_Wai = tf.assign(Wav,Wai)
        init_W_Wbi = tf.assign(Wbv,Wbi)

    else:
        Wbv = tf.cast(Wbi,dtype=tf.float64)
        Wav = tf.cast(Wai,dtype=tf.float64)
        
    Wb = tf.cast(tf.reshape(Wav,[M,1]),dtype=tf.float64) 
    Wa = tf.cast(tf.reshape(Wbv,[N,1]),dtype=tf.float64)         
    n_loop = tf.Variable(0, dtype=tf.int32)                         #循环计数
    Yout = tf.Variable(tf.zeros([batch_size,GlottisTime]),dtype=tf.float32)   #IIR输出
    n_loop_, X_, Yl_,Yout_,Wa_,Wb_ = tf.while_loop(cond, body, [n_loop, X, Yl,Yout,Wa,Wb]) #循环图
       
    ####### 训练方法 #################
       
    loss = tf.reduce_mean(tf.square(Yout_ - tf.to_float(Yd)))
    model_folder = '../model/'
    learning_rate = 10.001
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
    training_op = optimizer.minimize(loss)
        
    ####### sesson ###############
        
    init = tf.global_variables_initializer()
    iteration = 0

    #取数据
    Y_batch,Yl_batch = sample_get(wave_y_base,N) 
    
    #从文件导入Wa,Wb
    wa_data = np.zeros([N],dtype='float32')
    wb_data = np.zeros([M],dtype='float32') 
    
    wa_filepath = model_folder + 'wa.dat'
    wb_filepath = model_folder + 'wb.dat'
    if os.path.exists(os.path.abspath(wa_filepath)):
        LoadData(wa_filepath,wa_data)
    if os.path.exists(os.path.abspath(wb_filepath)):
        LoadData(wb_filepath,wb_data) 
    
    #从文件导入k
    glot_wav = np.zeros([GlottisTime],dtype='int16') 
    #m_filepath = model_folder + 'm74.snd'.format(GlottisTime) 
    #LoadWave(m_filepath,0,glot_wav,GlottisTime)  
    k_filepath = model_folder + 'k{}.snd'.format(GlottisTime)
    if os.path.exists(os.path.abspath(k_filepath)):
        LoadWave(k_filepath,0,glot_wav,GlottisTime)      
            
    k_data = glot_wav[::-1].astype(np.float32) #TF卷积核是颠倒的
    feed = {Yd:Y_batch,Yl:Yl_batch,Ki:k_data,Wai:wa_data,Wbi:wb_data}

    with tf.Session() as sess:
        sess.run(init) #全新训练的初始化       
        
        if param_mode=='Train' and param_type=='FIR':
            sess.run(init_K_Ki,feed_dict=feed)
        
        if param_mode=='Train' and param_type=='IIR':
            sess.run(init_W_Wai,feed_dict=feed)    
            sess.run(init_W_Wbi,feed_dict=feed)    
            
        while(param_mode == 'Train'):
            iteration += 1
            
            sess.run(training_op, feed_dict=feed)
                                                
            if iteration % 100 == 0:
                mse = loss.eval(feed_dict=feed)
                print(iteration, "\tMSE:", mse)
            
            if iteration % 2000 == 0:
                #保存Wa,Wb到文件
                if param_mode=='Train' and param_type=='IIR' :
                    wav = sess.run(Wav, feed_dict=feed).astype(np.float32)
                    wbv = sess.run(Wbv, feed_dict=feed).astype(np.float32)
                    SaveData(wa_filepath,wav)
                    SaveData(wb_filepath,wbv)
                    
                #保存k到文件
                if param_mode=='Train' and param_type=='FIR' :
                    k = sess.run(K, feed_dict=feed)
                    glot_wav = k[::-1].astype(np.int16)
                    SaveWave(k_filepath,glot_wav,1,False)
                
                print("\tmodel saved.")
                iteration = 0

        ######## 预测 ####################
        res = sess.run([n_loop_, X_, Yl_,Yout_,Wa_,Wb_],feed_dict=feed)
        wave_pred =  res[3][0]        
           
        #波形显示   
        show_one(glot_wav,'b') #声门波
        show_one(wave_y_base,'r')#目标元音
        show_one(wave_pred,'y') #元音预测
            
        # 增益递减 fooplot.com 1-x/1.1^(60-x)
        rate = np.zeros([30],dtype='float32')
        for n in range(30):
            rate[n] = 1-n/(1.1**(60-n))
            if (rate[n] < 0):rate[n] = 0
        
        # 连续递减保存 
        bAppend = False
        for n in range(30):                        
            wave_pred =  res[3][0] *rate[n]                                
            SaveWave('../data/gen_{}_{}.snd'.format(param_voice,GlottisTime),wave_pred,5,bAppend)
            bAppend = True
        


if __name__ == '__main__':
    #mode = 'Train'
    mode = 'Test'

    #type = 'IIR'
    type = 'FIR'

    voice = 'a'
    
    main(mode,type,voice)      
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    