#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Author: C0ffymachyne
License: GPLv3
Version: 1.0.0

Description:
    various plotting methods for debugging and visualization
"""

import torch
from io import BytesIO
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

def _figure_to_image(figure,format='png',dpi=96):
    buf = BytesIO()
    figure.savefig(buf, format=format, dpi=96)
    plt.close(figure)

    buf.seek(0)
    waveform_image = Image.open(buf).convert("RGB")

    return waveform_image

def get_wave(waveform : torch.Tensor, sample_rate : int , title : str ="Waveform", xlim=None, ylim=None):
    waveform = waveform.numpy()
    
    num_channels, num_frames = waveform.shape
    time_axis = torch.arange(0, num_frames) / sample_rate
    
    figure, axes = plt.subplots(num_channels, 1)
    if num_channels == 1:
        axes = [axes]
    for c in range(num_channels):
        axes[c].plot(time_axis, waveform[c], linewidth=1)
        axes[c].grid(True)
    
    if num_channels > 1:
        axes[c].set_ylabel(f'Channel {c+1}')
    if xlim:
        axes[c].set_xlim(xlim)
    if ylim:
        axes[c].set_ylim(ylim)
    
    figure.suptitle(title)

    waveform_image = _figure_to_image(figure)

    return waveform_image


def get_spectogram(waveform : torch.Tensor, sample_rate : int , title="Spectrogram", xlim=None): 
    waveform = waveform.numpy()
    num_channels, num_frames = waveform.shape
    time_axis = torch.arange(0, num_frames) / sample_rate
    
    figure, axes = plt.subplots(num_channels, 1)
    if num_channels == 1:
        axes = [axes]
    for c in range(num_channels):
        axes[c].specgram(waveform[c], Fs=sample_rate)
    if num_channels > 1:
        axes[c].set_ylabel(f'Channel {c+1}')
    if xlim:
        axes[c].set_xlim(xlim)
    figure.suptitle(title)

    waveform_image = _figure_to_image(figure)

    return waveform_image

def save_image(filepath : str, image : Image):
    image.save(filepath, format="JPEG", quality=95)