import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import fft
import pandas as pd
from scipy.fftpack import fft, ifft, fftshift, rfftfreq, rfft, irfft

import pandas as pd
import pywt
import torch
import plotly.graph_objects as go
import torchvision
from torchvision import transforms, datasets
import torch.nn as nn
import torch.nn.functional as F

def FFT(P):
    """ Функция реализует быстрое преобразование Фурье
    
    P - список сигналов"""
    
    n = len(P)
    if n == 1:
        return P
    else:
        w = np.exp((2*np.pi*1j)/n)
        Pe = [P[2*k] for k in range(int(n/2))]
        P0 = [P[2*k + 1] for k in range(int(n/2))]
        ye = FFT(Pe)
        y0 = FFT(P0)
        y = [0]*n
        for j in range(int(n/2)):
            y[j] = ye[j] + w**j*y0[j]
            y[j + int(n/2)] = ye[j] - w**j*y0[j]
        return y
        
def coefficient(sample, sample_name):
    
    """Функция считает коэфициент несинусоидальности
    sample - список сигналов (выборка)
    sample_name - название выборки"""
    
    peaks = np.abs(rfft([i[1] for i in sample]))
    index = np.argmax(peaks)
    span = []
    pr = (sample.shape[0] - 1)/(sample[-1][0] - sample[0][0])
    for i in range(1, sample.shape[0]//2):
        x = rfftfreq(sample.shape[0], 1/pr)[i]
        f = np.abs(rfft([item[1] for item in sample]))[i]
        span.append([x, f])
    Si = 0
    intr = span[2][0] - span[0][0]
    S0 = ((peaks[index-1] + peaks[index+1])/2 + peaks[index])*intr
    for i in range(1, peaks.size):
        Si += (peaks[i-1] + peaks[i])/2*intr
    Si -= S0
    print(f' coefficient {sample_name} =  {Si/S0}')
    
    
def plot_spectrum(sample, sample_name):
    
    """Функция считает коэфициент несинусоидальности
    sample - список сигналов (выборка)
    sample_name - название выборки"""
    
    # spectrum
    span = []
    pr = (sample.shape[0] - 1)/(sample[-1][0] - sample[0][0])
    for i in range(1, sample.shape[0]//2):
        x = rfftfreq(sample.shape[0], 1/pr)[i]
        f = np.abs(rfft([item[1] for item in sample]))[i]
        span.append([x, f])
    fig, axs = plt.subplots(1, 1, figsize=(9, 4))
    plt.plot([i[0] for i in span], [i[1] for i in span])
    plt.grid(True)
    plt.xlabel('Frequency')
    plt.ylabel('y_points')
    plt.title(f'Spectrum {sample_name}')
    plt.show()
    
    # beggining function
    fig, axs = plt.subplots(1, 1, figsize=(9, 4))
    plt.plot([el[0] for el in df], [el[1] for el in df])
    plt.grid(True)
    plt.title(f" Original function {sample_name}")
    plt.xlabel('x_points')
    plt.ylabel("y_points")
    
    # restart function
    fig, axs = plt.subplots(1, 1, figsize=(9, 4))
    peaks = list(np.abs(rfft([item[1] for item in sample])))
    rest = irfft(peaks)
    plt.plot(rest)
    plt.grid(True)
    plt.title(f"Restored function {sample_name}")
    plt.xlabel("x_points")
    plt.ylabel("y_points")
    
def variation(sample, sample_name):
    
    """Функция считает отклонение реальных данных от данных, получившихся в результате преобразование Фурье
    sample - список сигналов (выборка)
    sample_name - название выборки"""
    
    peaks = list(np.abs(rfft([item[1] for item in sample])))
    exmp = irfft(peaks)
    y = [line[1] for line in sample]
    var = sum([(y[i]-exmp[i])**2 for i in range(len(exmp))])
    print(f'variation {sample_name} : {var}')

def restart(lst, name): 
    
    """Функция восстановливает функции по спектрограмме
    lst - список сигналов"""

    period = (len(lst)-1)/(lst[-1][0] - lst[0][0])
    spektor = []
    for i in range(1, len(lst)//2):
        y = np.abs(rfft([el[1] for el in lst]))[i]
        x = rfftfreq(len(lst), 1/period)[i]
        spektor.append([x,y])
    peaks_n = np.abs(rfft([el[1] for el in lst]))
    peaks = list(peaks_n)
    part = len(lst)//10
    part_sp = len(spektor)//10
    fig, axs = plt.subplots(10, 3, figsize=(15, 30))
    plt.subplots_adjust(wspace=2, hspace=0.7)
    def plot_restart(i, n, n_sp):
       
        """Функция строит графики восстановленной функции по спектрограмме
        i - итерация, необходимая для определения длины отрезанного сигнала
        n - граница списка сигналов
        n_sp - граница списка спектров"""
        
        axs[i-1, 0].plot([el[0] for el in lst], [el[1] for el in lst])
        axs[i-1, 0].plot([el[0] for el in lst[:n]], [el[1] for el in lst[:n]], 'r')
        axs[i-1, 0].grid(True)
        axs[i-1, 0].set_title(f"Original function {name}")
        axs[i-1, 0].set_xlabel(f"x")
        axs[i-1, 0].set_ylabel(f"y")

        axs[i-1, 1].plot([el[0] for el in spektor[:n_sp]], [el[1] for el in spektor[:n_sp]])
        axs[i-1, 1].grid(True)
        axs[i-1, 1].set_title(f"Spectrum {name} (cup part {(i-1)*10}%)")
        axs[i-1, 1].set_xlabel('Frequency (Гц)')
        axs[i-1, 1].set_ylabel('Points')

        recover = irfft(peaks[:n])
        axs[i-1, 2].plot(recover)
        axs[i-1, 2].grid(True)
        axs[i-1, 2].set_title(f"Restart function {name}")
        axs[i-1, 2].set_xlabel(f"x")
        axs[i-1, 2].set_ylabel(f"y")
    
    for i in range(1,11):
        plot_restart(i, part*(11-i), part_sp*(11-i))
def descrete_decomposition(num_levels, y, name_sample):
    
    """Функция выводит Вейвлет-декомпозиции до уровня n
        num_levels - количество уровней декомпозиции
        y - список сигналов
        name_sample - названия списка сигналов"""
    
    for wev in ['haar', 'db2']:
        x = np.append(0, np.array(pd.read_csv(name_sample))[:, :1])
        plt.plot(x, y, 'orange')
        plt.grid(True)
        plt.title(f" {name_sample} : Исходная функция - {wev}")
        plt.show()
        for i in range(1, num_levels+1):
            # коэфициенты аппроксимации (выход фильтра нижних частот)
            coeffs = pywt.wavedec(y, wev, level=i)
            plt.subplot(2,1,1)
            plt.title(f'{name_sample} : Вейвлет - {wev}, декомпозиция {i}-ого уровня')
            plt.plot(coeffs[0], 'b', linewidth=2, label = f'cA, level={i}')
            plt.grid()
            plt.legend(loc = 'best')
            plt.show()
            # коэфициенты детализации (выход фильтра высоких частот)
            plt.subplot(2,1,2)
            plt.plot(coeffs[1], 'r', linewidth=2, label = f'cD, level={i}')
            plt.grid()
            plt.legend(loc = 'best')
            plt.show()
def count_level(func, wavelet):
    
    """Функция возвращает массивы аппроксимации и детализации 
    func - список сигналов
    wavelet - название вейвлета"""
    
    a, b = pywt.cwt(func, np.arange(1,4), wavelet)
    return a, b

def cont_decomposition(num_levels, y, name_sample):
    
    """Функция выводит Вейвлет-декомпозиции до уровня n
    num_levels - количество уровней декомпозиции
    y - список сигналов
    name_sample - названия списка сигналов"""
    
    for wev in ['gaus1', 'mexh']:
        x = np.append(0, np.array(pd.read_csv(name_sample))[:, :1])
        plt.plot(x, y, 'orange')
        plt.grid(True)
        plt.title(f" {name_sample} : Исходная функция - {wev}")
        plt.show()
        for i in range(1, num_levels+1):
            wavelet = pywt.ContinuousWavelet(wev, level = i)
            coef, freqs = count_level(y, wavelet)
            if i == 1:
                plt.subplot(2,1,1)
                plt.title(f' {name_sample}: Вейвлет - {wev}, декомпозиция {i}-ого уровня')
                plt.plot(coef[0], label = f' level={i}')
                plt.grid()
                plt.legend(loc = 'best')
                plt.show()
            else:
                for k in range(2, i+1):
                    coef, freqs = count_level(coef[0], wavelet)
                plt.subplot(2,1,1)
                plt.title(f' {name_sample}: Вейвлет - {wev}, декомпозиция {i}-ого уровня')
                plt.plot(coef[0], label = f' level={i}')
                plt.grid()
                plt.legend(loc = 'best')
                plt.show()
class Model(nn.Module):
    def __init__(self, in_feats, h_feats, out_feats):        
        super(Model, self).__init__()     # конструктор предка с этим именем
         
        self.fc1 = nn.Linear(in_feats, h_feats)             # создаём параметры модели
        self.fc2 = nn.Linear(h_feats, out_feats)             # в полносвязных слоях
          
    def forward(self, h):                        # задаётся прямой проход
        h = self.fc1(h)                          # выход первого слоя
        h = nn.Sigmoid()(h)         # пропускаем через Sigmoid
        h = self.fc2(h)
                        # выход второго слоя     
        return h
def Approximation(features, y, n_epochs, in_feats, h_feats, out_feats, N):
    
    """Функция аппроксимирует данную функцию
    features - значение аргумента функции
    y - список желаемых значений
    n_epoches - названия списка сигналов
    in_feats - количество входных слоев
    h_feats - количество скрытых слоев
    out_feats - количество выходных слоев
    N - число, определяющее границу выборки"""
    
    x = np.linspace(0, N, N)
    x.shape = 100, 1

    model = Model(in_feats, h_feats, out_feats) 
    optimizer = torch.optim.Adam(model.parameters(), lr=0.4)
    criterion = torch.nn.MSELoss()

    for epoch in range(n_epochs):

        optimizer.zero_grad() 
        # Forward
        pred = model(features)
        loss = criterion(pred.view(-1, 1), y2.view(-1, 1))

        # Вычисляем потери
        loss.backward()
        optimizer.step()

        if epoch % 500 == 0:
            print('In epoch {}, loss: {}'.format(epoch, loss))
            
    # рисуем нашу исходную функцию
    fig = go.Figure()
    fig.add_trace(go.Scatter(x = features.flatten().numpy(),
                                 y = y2.flatten().numpy(),
                                 mode = 'markers',
                                name = 'input_function'))

    # Генерируем  прогнозы для равномерно распределенных значений x между min и max
    min_x = min(list(features.numpy()))
    max_x = max(list(features.numpy()))
 
    z = torch.from_numpy(np.linspace(min_x, max_x, num = len(features))).reshape(-1, 1).float()
    w = model(z)

    # рисуем предсказанную функцию
    fig.add_trace(go.Scatter(x = z.flatten().numpy(),
                                 y = w.flatten().detach().numpy(),
                                 mode = 'lines',
                                name = 'predicted'))

    fig.show()
    
