import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from matplotlib.ticker import FixedLocator

f = 1                   # fréquence du signal (à modifier si besoin)

def sinus(t): 
    return np.sin(2. * np.pi * f * t)

def triangle(t):
    return signal.sawtooth(2. * np.pi * f * t,0.5)

def dent_scie(t):
    return signal.sawtooth(2. * np.pi * f * t,1)

def carre(t):
    return signal.square(2. * np.pi * f * t,.5)  # modifier .5 pour changer rapport cyclique

########


def sampling(s,F_ech,T_tot):
    N = int(T_tot * F_ech)                 # nombre points échantillon
    te = np.linspace(0., T_tot, N+1)       # liste des instants d'échantillonage
    tp = np.linspace(0,T_tot,1000)         # abscisse du graphe du signal initial
    plt.close()
    plt.plot(te, s(te), 'or', markersize = 4, label = u"Signal echantillonne")
    plt.plot(te, s(te), 'r--',linewidth=2)
    plt.plot(tp, s(tp), 'b', markersize = 4, label = u"Signal reel")
    plt.grid()
    plt.xlabel("$t$")
    plt.ylabel("$s(t),s_n$")
    plt.legend(loc="upper right")
    plt.title("$F_\mathrm{éch}=$"+str(F_ech)+" Hz")
    plt.show()
    return

########


def EB(s,F_ech,T_tot):
    N = int(T_tot * F_ech)                   # nombre de points échantillon
    T_ech = 1/F_ech
    time_ech = [0]
    Ueb = [s(0)]
    for i in range(N):
        time_ech.append(i*T_ech)
        time_ech.append(i*T_ech)
        Ueb.append(Ueb[-1])
        Ueb.append(s(i*T_ech))
    time_ech.append(N*T_ech)
    Ueb.append(s((N-1)*T_ech))
    tp = np.linspace(0,T_tot,1000)
    plt.close()
    plt.plot(time_ech,Ueb, label = u"Signal echantillonne")
    plt.plot(tp,s(tp), label = u"Signal reel")
    plt.xlabel("$t$")
    plt.legend(loc='upper right')
    plt.title("$F_\mathrm{éch}=$"+str(F_ech)+" Hz")
    plt.show()
    return

########


def EBQ(s,F_ech,T_tot,Cal,q):
    N = int(T_tot * F_ech)                 # nombre de points échantillon
    T_ech = 1/F_ech
    tp = np.linspace(0,T_tot,1000)
    time_ech = [0]                         # abscisses du signal bloqué
    n = 2**q                               # nombre valeurs quantifiées
    pas = 2*Cal/(n-1)                      # pas quantification
    valQ = [-Cal]
    for i in range(1,n):
        valQ.append(-Cal+i*pas)            # liste des valeurs accessibles par quantification 
    time_ech = [0]
    Ueb = [s(0)]   
    for i in range(N+1):                    
        time_ech.append(i*T_ech)
        time_ech.append(i*T_ech)
        Ueb.append(Ueb[-1])
        Ueb.append(s(i*T_ech))
    UebQ = []
    for x in Ueb:
        indice = int((x+Cal)/pas)
        if x < -Cal:
            UebQ.append(-Cal)
        elif x > Cal:
            UebQ.append(Cal)
        elif np.abs(x-valQ[indice]) <= pas/2:
            UebQ.append(valQ[indice])
        else:
            UebQ.append(valQ[indice+1])
    plt.close()
    fig, ax = plt.subplots()
    ax.grid(axis='y')
    while len(valQ)>40:                 # pour éviter d'avoir plus de 40 valeurs de graduation en y   
        valQ = valQ[::2]
    major_locator = FixedLocator(valQ)
    ax.yaxis.set_major_locator(major_locator)
    plt.ylim(-Cal,Cal)
    ax.plot(time_ech,UebQ, label = u"Signal echantillonne")
    ax.plot(tp,s(tp), label = u"Signal reel")
    plt.xlabel("$t$")
    plt.legend()
    plt.title("$F_\mathrm{éch}=$"+str(F_ech)+" Hz"+" et échantillonnage sur "+str(q)+" bits")
    plt.show()
    return