#!/usr/bin/env python3

__doc__ = """
Phase and Amplitude distributions for a phasor in the presence of noise.

Author: Harm Schoorlemmer
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import scipy.stats as stat
from scipy import special

from lib.util import MethodMappingProxy as MethodProxy
from lib.ft_plot import axis_pi_ticker

def expectation(x,pdfx):
    dx = x[1]-x[0]
    return np.sum(x*pdfx*dx)

def variance(x,pdfx):
    mu = expectation(x,pdfx)
    dx = x[1]-x[0]
    return np.sum((x**2*pdfx*dx))-mu**2

def phase_distribution(theta,sigma,s):
    theta = np.asarray(theta)
    ct = np.cos(theta)
    st = np.sin(theta)
    k=s/sigma
    pipi=2*np.pi
    return (np.exp(-k**2/2)/pipi) + (
        (pipi**-0.5)*k*np.exp(-(k*st)**2/2)) * (
        (1.+special.erf(k*ct*2**-0.5))*ct/2)

def phase_distribution_gauss(theta,sigma,s):
    theta = np.asarray(theta)
    k=s/sigma
    return (2*np.pi)**-0.5*k*np.exp(-(k*theta)**2/2)

def amplitude_distribution(a,sigma,s):
    a = np.asarray(a)
    k = s/sigma
    return stat.rice.pdf(a,s/sigma,scale=sigma)

def amplitude_distribution_gauss(a,sigma,s):
    a = np.asarray(a)
    k=s/sigma
    return (2*np.pi)**-0.5*np.exp(-((a-s)/sigma)**2/2)


figsize = (8,6)
if True:
        from matplotlib import rcParams
        #rcParams["text.usetex"] = True
        rcParams["font.family"] = "serif"
        plt.rc('lines',lw=2)
        if True:# small
            figsize = (6, 4)
            rcParams["font.size"] = "15" # 15 at 6,4 looks fine
        elif True: # large
            figsize = (9, 6)
            rcParams["font.size"] = "16" # 15 at 9,6 looks fine
        rcParams["grid.linestyle"] = 'dotted'
        rcParams["figure.figsize"] = figsize


signal_max= 4
amp_max = signal_max*2
thetas = np.linspace(-np.pi,np.pi,500)
amplitudes = np.linspace(0,amp_max,500)
signals = np.linspace(0,signal_max,5)
sigma = 1


## figure 1
if True:
    fig, ax = plt.subplots(1,2,figsize=(2*figsize[0],figsize[1]))
    _fig1, _ax0 =plt.subplots(1,1)
    _fig2, _ax1 =plt.subplots(1,1)

    ax0 = MethodProxy(ax[0], _ax0)
    ax1 = MethodProxy(ax[1], _ax1)
    for s in signals:
        pdfs_label='s = '+str(int(s))
        phase_vals= phase_distribution(thetas,sigma,s)
        amp_vals= amplitude_distribution(amplitudes,sigma,s)
        phase_vals_g = phase_distribution_gauss(thetas,sigma,s)

        ax0.plot(amplitudes,amp_vals, label=pdfs_label)
        ax1.plot(thetas,phase_vals,   label=pdfs_label)
    ax0.legend()
    ax0.set_xlabel(r'Amplitude $a$')
    ax0.set_ylabel(r'$p(a)$')

    ax1.set_xlabel(r'Phase $\varphi$')
    ax1.set_ylabel(r'$p(\varphi)$')
    _ax1.legend()# only in the separate figure
    [ axis_pi_ticker(ax.xaxis, major_divider=3) for ax in ax1.elements ]

    for a in [ax0, ax1]:
        a.grid()
    MethodProxy(fig, _fig1, _fig2).tight_layout()

    fig.savefig('pdfs.pdf')
    _fig1.savefig('pdfs-amplitudes.pdf')
    _fig2.savefig('pdfs-phases.pdf')

    plt.close(_fig1)
    plt.close(_fig2)

## figure 2
amplitudes = np.linspace(0,amp_max*5,500)
signals = np.linspace(0.1,signal_max*5,101)
if True:
    V_theta = [variance(thetas,phase_distribution(thetas,sigma,s)) for s in signals ]
    E_theta=[expectation(thetas,phase_distribution(thetas,sigma,s)) for s in signals ]
    V_theta_g = [variance(thetas,phase_distribution_gauss(thetas,sigma,s)) for s in signals ]
    E_theta_g=[expectation(thetas,phase_distribution_gauss(thetas,sigma,s)) for s in signals ]
    V_a = [variance(amplitudes,amplitude_distribution(amplitudes,sigma,s)) for s in signals ]
    E_a=[expectation(amplitudes,amplitude_distribution(amplitudes,sigma,s)) for s in signals ]
    V_a_g = [variance(amplitudes,amplitude_distribution_gauss(amplitudes,sigma,s)) for s in signals ]
    E_a_g=[expectation(amplitudes,amplitude_distribution_gauss(amplitudes,sigma,s)) for s in signals ]

    fig2, _ax2 = plt.subplots(2,2,figsize=(2*figsize[0],2*figsize[1]))
    ax2 = fig2.get_axes()
    if True:
        _figs = []
        _axs = []
        for i, ax in enumerate(ax2):
            _f, _a = plt.subplots(1,1)
            _figs.append(_f)
            _axs.append(_a)

            ax2[i] = MethodProxy(ax, _a)

    ax2[0].plot(signals,E_a,label='$p(a)$')
    ax2[0].plot(signals,E_a_g,ls='dashed',label='Gaussian approx.')
    ax2[0].set_xscale('log')
    ax2[0].set_yscale('log')
    ax2[0].set_ylabel('$\mu_a$')

    ax2[1].plot(signals,V_a,label='$p(a)$')
    ax2[1].plot(signals,V_a_g,ls='dashed',label='Gaussian approx.')
    ax2[1].set_xscale('log')
    ax2[1].set_ylabel('$\sigma_a^2$')

    ax2[2].plot(signals,E_theta,label=r'$p(\varphi)$')
    ax2[2].plot(signals,E_theta_g,ls='dashed',label='Gaussian approx.')
    ax2[2].set_xscale('log')
    ax2[2].set_ylim(-1.1,1.1)
    ax2[2].set_ylabel(r'$\mu_\varphi$')

    ax2[3].plot(signals,V_theta,label=r'$p(\varphi)$')
    ax2[3].plot(signals,V_theta_g,ls='dashed',label='Gaussian approx.')
    ax2[3].set_xscale('log')
    ax2[3].set_yscale('log')
    ax2[3].set_ylabel(r'$\sigma_\varphi^2$')
    for a in ax2:
        a.grid(which='both')
        a.set_xlabel(r'$s/\sigma$')
        a.legend()

    fig2.tight_layout()
    fig2.savefig('expectation_variance.pdf')
    for i, _f in enumerate(_figs):
        fnames = [
                'amplitude_mean',
                'amplitude_sigma',
                'phase_mean',
                'phase_sigma',
            ][i]

        _f.tight_layout()
        _f.savefig(fnames+'.pdf')
        plt.close(_f)

## figure 3, beacon timing accuracy
if True:
    fig3, ax3 = plt.subplots(1,1)
    ax3 = fig3.get_axes()
    sigma_t = [variance(thetas,phase_distribution(thetas,sigma,s)) for s in signals ]
    lfs=np.linspace(np.log10(50.),4,1)
    for lf in lfs:
        freq = (10**lf)*1e6
        sigma_t = np.asarray(sigma_t)**0.5/(2*np.pi*freq)
        ax3[0].plot(signals,sigma_t/1e-9,'o-')
        ax3[0].set_ylim(0,2.5)
        ax3[0].set_xlabel(r'$s/\sigma$')
        ax3[0].set_ylabel(r'$\Delta t$(ns)')

    fig3.tight_layout()
    fig3.savefig('timing_accuracy.pdf')

plt.show()