import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib.dates as mdates


def plot_annual_data(csv_path, index, file_name, title_img, caption_img, embed_mode=False):
    data_base = pd.read_csv(csv_path)
    data_base.columns = ["year", "prec", "tmax", "tmin", "tmean"]
    data_base[index] = data_base[index].replace(-99.9, None)
    data_base["date"] = pd.to_datetime(data_base["year"].astype(str) + "-01-01")

    sns.set_style("white")
    fig, ax = plt.subplots(figsize=(12, 1.5))
    cmap = sns.color_palette("RdBu_r", as_cmap=True)

    scatter = ax.scatter(
        data_base["date"],
        [1] * len(data_base),
        c=data_base[index],
        cmap=cmap,
        marker="s",
        s=800
    )

    ax.xaxis.set_major_locator(mdates.YearLocator(5))
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y"))
    
    ax.set_xlim(data_base["date"].min() - pd.Timedelta(days=200), 
                data_base["date"].max() + pd.Timedelta(days=200))
    
    ax.set_ylim(0.9, 1.1) 
    ax.set_yticks([])
    
    ax.set_title(title_img, fontsize=13, fontweight="bold", pad=10)
    ax.grid(False)
    
    fig.colorbar(scatter, ax=ax, label=f"Temp ({index})", pad=0.01, aspect=10)
    
    ax.text(1.0, -0.4, caption_img, transform=ax.transAxes, fontsize=9, ha="right")

    if embed_mode:
        return fig
    else:
        fig.savefig(file_name, dpi=300, bbox_inches='tight', pad_inches=0.02)
        plt.show()
        plt.close(fig)
        return None


def plot_quarterly_data(csv_path, index, file_name, title_img, caption_img, embed_mode=False):
    data_base = pd.read_csv(csv_path)
    data_base.columns = ["year", "quarter", "prec", "tmax", "tmin", "tmean"]
    data_base["year"] = data_base["year"].astype(int)
    data_base[index] = data_base[index].replace(-99.9, None)

    pivot_table = data_base.pivot(index="quarter", columns="year", values=index)
    pivot_table = pivot_table.sort_index(ascending=False) # T4 em cima, T1 embaixo

    sns.set_style(style="white")
    fig, ax = plt.subplots(figsize=(16, 4)) # Largura maior, altura menor
    cmap = sns.color_palette("RdBu_r", as_cmap=True)

    # O HEATMAP preenche todo o espaço automaticamente
    sns.heatmap(
        pivot_table, 
        cmap=cmap, 
        ax=ax, 
        cbar_kws={'label': f"Temperatura ({index})", 'pad': 0.01},
        linewidths=0, # 0 para colar os quadrados, ou 0.5 para uma linha fina entre eles
        xticklabels=True, 
        yticklabels=["T4", "T3", "T2", "T1"]
    )

    ax.set_title(title_img, fontsize=16, fontweight="bold", pad=20)
    ax.set_xlabel("Ano")
    ax.set_ylabel("Trimestre")
    
    # Legenda (Caption)
    ax.text(1.0, -0.25, caption_img, transform=ax.transAxes, fontsize=10, ha="right")

    if embed_mode:
        return fig
    else:
        fig.savefig(file_name, dpi=300, bbox_inches='tight')
        plt.show()
        plt.close(fig)


def plot_monthly_data(csv_path, index, file_name, title_img, caption_img, embed_mode=False):
    data_base = pd.read_csv(csv_path)
    data_base.columns = ["year", "month", "prec", "tmax", "tmin", "tmean"]
    data_base[index] = data_base[index].replace(-99.9, None)

    sns.set_style(style="white")
    fig, ax = plt.subplots(figsize=(14, 5))
    cmap = sns.color_palette("RdBu_r", as_cmap=True)

    scatter = ax.scatter(
        data_base["year"],
        data_base["month"],
        c=data_base[index],
        cmap=cmap,
        marker="s",
        s=450
    )

    ax.invert_yaxis() 
    ax.set_xlim(data_base["year"].min() - 0.5, data_base["year"].max() + 0.5)
    ax.tick_params(axis='x', rotation=90)
    ax.set_yticks(range(1, 13), labels=[
        "Jan", "Fev", "Mar", "Abr", "Mai", "Jun",
        "Jul", "Ago", "Set", "Out", "Nov", "Dez"
    ])
    ax.set_title(title_img, fontsize=14, fontweight="bold")
    ax.set_xlabel("Ano")
    ax.set_ylabel("Mês")
    ax.grid(False) 
    cbar = fig.colorbar(scatter, ax=ax, label=f"Temperature ({index})", aspect=30, pad=0.02)
    ax.text(1.0, -0.15, caption_img, transform=ax.transAxes, fontsize=10, ha="right")
    fig.tight_layout(pad=1.0)

    if embed_mode:
        return fig
    else:
        fig.savefig(file_name, dpi=300, bbox_inches='tight')
        plt.show()
        plt.close(fig)
        return None
