from io import BytesIO
from quart import Blueprint, request, send_file
from matplotlib import pyplot as plt

from modules.cosmic.gravityw_multiday import GravityMultidayPlot
from modules.cosmic.gravityw_perday import CosmicGravitywPlot
from modules.cosmic.planetw_daily import cosmic_planetw_daily_plot
from quart.utils import run_sync


cosmic_module = Blueprint("Cosmic", __name__)


@cosmic_module.route('/metadata')
def get_meta():
    return []


@cosmic_module.route('/render/planet_wave/daily')
@cosmic_module.route('/temp_render')
async def render():
    T = request.args.get("T", 16)
    await run_sync(cosmic_planetw_daily_plot)(T=int(T))
    buf = BytesIO()

    plt.savefig(buf, format="png")
    buf.seek(0)
    return await send_file(buf, mimetype="image/png")


@cosmic_module.route('/render/gravity_wave/perday')
@cosmic_module.route('/render/single')
async def single_render():
    year = request.args.get("year", 2008)
    day = request.args.get("day", 1)
    mode = request.args.get("mode", "mean_ktemp_Nz")

    p: CosmicGravitywPlot = await run_sync(CosmicGravitywPlot)(year=int(year), day=int(day))
    if mode == "mean_ktemp_Nz":
        await run_sync(p.plot_results_mean_ktemp_Nz)()
    elif mode == "mean_ktemp_Ptz":
        await run_sync(p.plot_results_mean_ktemp_Ptz)()
    else:
        raise ValueError("Invalid mode")

    buf = BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)

    return await send_file(buf, mimetype="image/png")


@cosmic_module.route("/render/gravity_wave/multiday")
async def multiday_render():
    year = request.args.get("year", 2008)
    start_day = request.args.get("start_day", 1)
    end_day = request.args.get("end_day", 204)
    mode = request.args.get("mode", "位温分布")

    p: GravityMultidayPlot = await run_sync(GravityMultidayPlot)(year=int(year),
                                                                 day_range=(start_day, end_day))
    if mode == "布伦特-维萨拉频率分布":
        await run_sync(p.plot_heatmap_tempNz)()
    elif mode == "位温分布":
        await run_sync(p.plot_heatmap_tempPtz)()
    elif mode == "每月浮力频率变化趋势":
        await run_sync(p.plot_floatage_trend)()
    elif mode == "每月平均重力势能的折线图":
        await run_sync(p.plot_monthly_energy)()
    else:
        raise ValueError("Invalid mode")

    buf = BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)

    return await send_file(buf, mimetype="image/png")