from sparkstart import scon, spark
from pyspark.sql import SparkSession
import time
import matplotlib.pyplot as plt

HDFSPATH = "hdfs://193.174.205.250:54310/"

def read_parquets(spark: SparkSession):
	stations_path = HDFSPATH + "home/heiserervalentin/german_stations.parquet"
	products_path = HDFSPATH + "home/heiserervalentin/german_stations_data.parquet"

	stations_df = spark.read.parquet(stations_path)
	stations_df.createOrReplaceTempView("german_stations")

	products_df = spark.read.parquet(products_path)
	products_df.createOrReplaceTempView("german_stations_data")

	stations_df.cache()
	products_df.cache()


def plot_all_stations(spark: SparkSession):
	q = "SELECT geo_laenge AS lon, geo_breite AS lat FROM german_stations WHERE geo_laenge IS NOT NULL AND geo_breite IS NOT NULL"
	df = spark.sql(q)

	pdf = df.toPandas()
	plt.figure(figsize=(8, 6))
	plt.scatter(pdf.lon, pdf.lat, s=6, color='red', marker='.')
	plt.xlabel('Longitude')
	plt.ylabel('Latitude')
	plt.title('All Stations (locations)')
	plt.tight_layout()
	plt.show()


def duration_circle_size(spark: SparkSession):
	q = (
		"SELECT stationId, geo_laenge AS lon, geo_breite AS lat, "
		"(CAST(SUBSTR(bis_datum,1,4) AS INT) - CAST(SUBSTR(von_datum,1,4) AS INT)) AS duration_years "
		"FROM german_stations "
		"WHERE TRIM(von_datum)<>'' AND TRIM(bis_datum)<>''"
	)
	df = spark.sql(q)

	pdf = df.toPandas()

	pdf['duration_years'] = pdf['duration_years'].fillna(0).astype(int)
	sizes = (pdf['duration_years'].clip(lower=0) + 1) * 6

	plt.figure(figsize=(8, 6))
	plt.scatter(pdf.lon, pdf.lat, s=sizes, alpha=0.6, c=pdf['duration_years'], cmap='viridis')
	plt.colorbar(label='Duration (years)')
	plt.xlabel('Longitude')
	plt.ylabel('Latitude')
	plt.title('Stations with duration (years) as marker size')
	plt.tight_layout()
	plt.show()


def compute_daily_and_yearly_frosts(spark: SparkSession):
    q_daily_max = (
        "SELECT stationId, date, SUBSTR(CAST(date AS STRING),1,4) AS year, MAX(TT_TU) AS max_temp "
        "FROM german_stations_data "
        "WHERE TT_TU IS NOT NULL AND TT_TU > -50 AND TT_TU < 60 "
        "GROUP BY stationId, date"
    )
    daily_max = spark.sql(q_daily_max)
    daily_max.createOrReplaceTempView('daily_max')

    # mark a day as frost if max_temp < 0
    q_daily_frost = (
        "SELECT stationId, year, CASE WHEN max_temp < 0 THEN 1 ELSE 0 END AS is_frost "
        "FROM daily_max"
    )
    daily_frost = spark.sql(q_daily_frost)
    daily_frost.createOrReplaceTempView('daily_frost')

    # yearly frostdays per station
    q_station_year = (
        "SELECT stationId, year, SUM(is_frost) AS frost_days "
        "FROM daily_frost GROUP BY stationId, year"
    )
    station_year_frost = spark.sql(q_station_year)
    station_year_frost.createOrReplaceTempView('station_year_frost')


def frost_analysis(spark: SparkSession, year=2024, station_name_matches=('kempten',)):
    compute_daily_and_yearly_frosts(spark)

    q_hist = (
        f"SELECT frost_days, COUNT(*) AS station_count "
        f"FROM station_year_frost WHERE year = '{year}' GROUP BY frost_days ORDER BY frost_days"
    )
    hist_df = spark.sql(q_hist)

    hist_pdf = hist_df.toPandas()
    
    if hist_pdf.empty:
        print(f"No frost data found for year {year}. Trying to find available years...")
        q_all = "SELECT frost_days, COUNT(*) AS station_count FROM station_year_frost GROUP BY frost_days ORDER BY frost_days"
        hist_pdf = spark.sql(q_all).toPandas()
        if hist_pdf.empty:
            print("No frost data available at all. Check if TT_TU column contains valid temperature data.")
            return
        print(f"Found {len(hist_pdf)} frost day categories across all years")
    
    plt.figure(figsize=(8, 5))
    plt.bar(hist_pdf.frost_days, hist_pdf.station_count, color='steelblue')
    plt.xlabel('Number of Frost Days in year ' + str(year))
    plt.ylabel('Number of Stations')
    plt.title(f'Stations vs Frost Days ({year})')
    plt.tight_layout()
    plt.show()

    for name in station_name_matches:
        q_find = f"SELECT stationId, station_name FROM german_stations WHERE lower(station_name) LIKE '%{name.lower()}%'"
        ids_df = spark.sql(q_find)
        ids = ids_df.collect()
        if not ids:
            print(f"No stations found matching '{name}'")
            continue
        for r in ids:
            sid = r['stationId']
            sname = r['station_name']
            print(f"Analyzing stationId={sid} name={sname}")

            q_ts = (
                "SELECT year, frost_days, "
                "AVG(frost_days) OVER (PARTITION BY stationId ORDER BY CAST(year AS INT) RANGE BETWEEN 4 PRECEDING AND CURRENT ROW) AS avg_5, "
                "AVG(frost_days) OVER (PARTITION BY stationId ORDER BY CAST(year AS INT) RANGE BETWEEN 19 PRECEDING AND CURRENT ROW) AS avg_20 "
                f"FROM station_year_frost WHERE stationId = {sid} ORDER BY CAST(year AS INT)"
            )
            ts_df = spark.sql(q_ts)

            pdf = ts_df.toPandas()
            if pdf.empty:
                print(f"No yearly frost data for station {sid}")
                continue

            pdf['year'] = pdf['year'].astype(int)
            plt.figure(figsize=(10, 5))
            plt.plot(pdf.year, pdf.frost_days, label='Frostdays (year)', marker='o')
            plt.plot(pdf.year, pdf.avg_5, label='5-year avg', linestyle='--')
            plt.plot(pdf.year, pdf.avg_20, label='20-year avg', linestyle=':')
            plt.xlabel('Year')
            plt.ylabel('Frost Days')
            plt.title(f'Frost Days over Years for {sname} (station {sid})')
            plt.legend()
            plt.tight_layout()
            plt.show()


def height_frost_correlation(spark: SparkSession):
	compute_daily_and_yearly_frosts(spark)

	q_corr = (
		"SELECT syf.year AS year, corr(s.hoehe, syf.frost_days) AS height_frost_corr "
		"FROM station_year_frost syf JOIN german_stations s ON syf.stationId = s.stationId "
		"GROUP BY syf.year ORDER BY CAST(syf.year AS INT)"
	)
	
	corr_df = spark.sql(q_corr)

	corr_pdf = corr_df.toPandas()

	corr_pdf = corr_pdf.dropna(subset=['height_frost_corr'])
	if corr_pdf.empty:
		print("No non-NaN correlation values found.")
		return

	corr_pdf['year'] = corr_pdf['year'].astype(int)
	plt.figure(figsize=(10, 5))
	plt.bar(corr_pdf.year, corr_pdf.height_frost_corr, color='orange')
	plt.xlabel('Year')
	plt.ylabel('Correlation (height vs frostdays)')
	plt.title('Yearly correlation: station height vs number of frost days')
	plt.tight_layout()
	plt.show()


def main(scon, spark):
	read_parquets(spark)

	plot_all_stations(spark)

	duration_circle_size(spark)

	frost_analysis(spark, year=2024, station_name_matches=('kempten',))

	height_frost_correlation(spark)


if __name__ == '__main__':
	main(scon, spark)

