#!/usr/bin/env python3
import os
import sys
from datetime import datetime
import calendar
import time
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.patches import Polygon
import pandas as pd

marker = 'o'
markersize = 4
linestyle = '-'
linewidth = 2
colors = {
    'max': 'red',
    'min': 'blue',
    'fill': 'gray'
}

def read_file_values(path):
    temp = pressure = humidity = None
    temp_u = pressure_u = humidity_u = ''
    try:
        with open(path, 'r') as f:
            for line in f:
                if 'Temperature' in line:
                    try:
                        parts = line.split(':', 1)[1].strip().split()
                        temp = float(parts[0])
                        temp_u = ' '.join(parts[1:]) if len(parts) > 1 else ''
                    except:
                        pass
                elif 'Pressure' in line:
                    try:
                        parts = line.split(':', 1)[1].strip().split()
                        pressure = float(parts[0])
                        pressure_u = ' '.join(parts[1:]) if len(parts) > 1 else ''
                    except:
                        pass
                elif 'Humidity' in line:
                    try:
                        parts = line.split(':', 1)[1].strip().split()
                        humidity = float(parts[0])
                        humidity_u = ' '.join(parts[1:]) if len(parts) > 1 else ''
                    except:
                        pass
    except FileNotFoundError:
        return (None, '', None, '', None, '')
    return (temp, temp_u, pressure, pressure_u, humidity, humidity_u)

def read_measurements(data_folder):
    records = []
    year_folder = str(datetime.now().year)
    for root, dirs, files in os.walk(data_folder):
        for filename in files:
            if not filename.endswith('.txt'):
                continue
            if filename == year_folder:
                continue
            filepath = os.path.join(root, filename)
            try:
                timestamp = int(filename[:-4])
            except ValueError:
                print(f"Filename not a timestamp, skipping: {filename} (in {root})")
                continue
            file_time = datetime.fromtimestamp(timestamp)

            temp, temp_u, pressure, pressure_u, humidity, humidity_u = read_file_values(filepath)

            # fallback to year subfolder file if some values missing
            if (temp is None or pressure is None or humidity is None):
                alt_path = os.path.join(data_folder, year_folder, filename)
                if os.path.isfile(alt_path):
                    at, at_u, ap, ap_u, ah, ah_u = read_file_values(alt_path)
                    if temp is None and at is not None:
                        temp, temp_u = at, at_u
                    if pressure is None and ap is not None:
                        pressure, pressure_u = ap, ap_u
                    if humidity is None and ah is not None:
                        humidity, humidity_u = ah, ah_u

            if temp is not None and pressure is not None and humidity is not None:
                records.append({
                    'timestamp': file_time,
                    'date': file_time.date(),
                    'day': file_time.day,
                    'temperature': temp,
                    'temperature_unit': temp_u,
                    'pressure': pressure,
                    'pressure_unit': pressure_u,
                    'humidity': humidity,
                    'humidity_unit': humidity_u
                })
            else:
                print(f"Missing data in file: {filename} (in {root}). Skipping.")
    return pd.DataFrame(records)

def filter_month(df, month, year):
    if df.empty:
        return df
    return df[(df['timestamp'].dt.month == month) & (df['timestamp'].dt.year == year)]

def daily_min_max(df, col):
    g = df.groupby('date')[col].agg(['min', 'max']).reset_index()
    g['day'] = pd.to_datetime(g['date']).dt.day
    g = g.sort_values('day')
    return g

def plot_min_max(daily_df, col, out_folder, year, month, unit=''):
    if daily_df.empty:
        print(f"No data to plot for {col}")
        return

    last_day = calendar.monthrange(year, month)[1]

    fig, ax = plt.subplots(figsize=(12, 6))
    ax.plot(daily_df['day'], daily_df['min'],
            marker=marker, markersize=markersize,
            linestyle=linestyle, linewidth=linewidth,
            label='Min', color=colors['min'], zorder=3)
    ax.plot(daily_df['day'], daily_df['max'],
            marker=marker, markersize=markersize,
            linestyle=linestyle, linewidth=linewidth,
            label='Max', color=colors['max'], zorder=4)

    xs = list(daily_df['day'])
    ys_max = list(daily_df['max'])
    ys_min = list(daily_df['min'])

    if len(xs) >= 2:
        verts = [(x, y) for x, y in zip(xs, ys_max)] + [(x, y) for x, y in zip(xs[::-1], ys_min[::-1])]

        cmap = LinearSegmentedColormap.from_list('blue_red', [colors['min'], colors['max']])
        global_min = min(ys_min)
        global_max = max(ys_max)
        if global_min == global_max:
            global_min -= 0.5
            global_max += 0.5

        gradient = np.linspace(0, 1, 256).reshape(256, 1)
        gradient = np.repeat(gradient, 10, axis=1)

        im = ax.imshow(gradient, aspect='auto', cmap=cmap,
                       extent=(0.5, last_day + 0.5, global_min, global_max),
                       origin='lower', alpha=0.35, zorder=1)

        poly = Polygon(verts, closed=True, transform=ax.transData)
        im.set_clip_path(poly)

    ax.set_xlim(1, last_day)
    xticks = list(range(1, last_day + 1))
    xtick_labels = [f"{d}." for d in xticks]
    ax.set_xticks(xticks)
    ax.set_xticklabels(xtick_labels, rotation=0)

    ylabel = col.capitalize()
    if unit:
        ylabel += f' in {unit}'
    ax.set_ylabel(ylabel)
    ax.set_xlabel(f'Day of Month ({year}-{month:02d})')
    ax.set_title(f'{col.capitalize()} - Daily Min/Max Value for {year}-{month:02d}')
    ax.grid(True, alpha=0.25)
    ax.legend()
    fig.tight_layout()
    out_path = os.path.join(out_folder, f'{col}_daily_min_max.png')
    fig.savefig(out_path)
    plt.close(fig)
    print(f"Saved plot: {out_path}")

def ensure_folder(path):
    if not os.path.exists(path):
        os.makedirs(path, exist_ok=True)

def pick_unit(series_unit_col):
    if series_unit_col is None or series_unit_col.empty:
        return ''
    vals = series_unit_col.dropna().astype(str)
    vals = vals[vals != '']
    return vals.mode().iat[0] if not vals.empty else ''

def main():
    if len(sys.argv) < 2:
        print("Usage: python3 script.py <month_number (1-12)>")
        sys.exit(1)
    try:
        month = int(sys.argv[1])
        assert 1 <= month <= 12
    except:
        print("Month must be an integer between 1 and 12.")
        sys.exit(1)

    data_folder = 'temperature'
    if not os.path.isdir(data_folder):
        print(f"Data folder not found: {data_folder}")
        sys.exit(1)

    start_time = time.time()
    df = read_measurements(data_folder)
    if df.empty:
        print("No valid measurements found.")
        sys.exit(1)

    df['timestamp'] = pd.to_datetime(df['timestamp'])

    year = datetime.now().year
    df_month = filter_month(df, month, year)
    if df_month.empty:
        print(f"No measurements for {month}/{year}.")
        sys.exit(0)

    # pick most common unit per measurement type for the month
    units = {
        'temperature': pick_unit(df_month.get('temperature_unit')),
        'pressure': pick_unit(df_month.get('pressure_unit')),
        'humidity': pick_unit(df_month.get('humidity_unit'))
    }

    out_folder = f'plots_{year}_{month:02d}'
    ensure_folder(out_folder)

    for col in ['temperature', 'pressure', 'humidity']:
        daily = daily_min_max(df_month, col)
        plot_min_max(daily, col, out_folder, year, month, unit=units.get(col, ''))

    elapsed = time.time() - start_time
    print(f"Verarbeitung abgeschlossen in {elapsed:.2f}s. Plots in {out_folder}")

if __name__ == '__main__':
    main()