import pandas as pd
import numpy as np
import argparse

# Command-line arguments
parser = argparse.ArgumentParser(description="Quality control of precipitation observations.")
parser.add_argument("arquivo", help="Input file (e.g., Teste_201801.txt)")
parser.add_argument("saida_aceitas", help="Output file for accepted observations")
parser.add_argument("saida_rejeitadas", help="Output file for rejected observations")
args = parser.parse_args()

# Column names
colunas = [
    "ID", "Longitude", "Latitude",
    "Observado", "Estimado", "Climatologia", "Regiao"
]

# Read the input file
df = pd.read_csv(args.arquivo, delim_whitespace=True, header=None, names=colunas)

# Compute absolute errors
df["Erro_abs_sat"] = (df["Observado"] - df["Estimado"]).abs()
df["Erro_abs_clim"] = (df["Observado"] - df["Climatologia"]).abs()

# Quality control function by region
def controle_qualidade_regional(df, k_sat=1.0, k_clim=1.0):
    aceitas = []
    rejeitadas = []

    for regiao, grupo in df.groupby("Regiao"):
        # Median and MAD of the errors
        mediana_sat = grupo["Erro_abs_sat"].median()
        mad_sat = np.median(np.abs(grupo["Erro_abs_sat"] - mediana_sat))

        mediana_clim = grupo["Erro_abs_clim"].median()
        mad_clim = np.median(np.abs(grupo["Erro_abs_clim"] - mediana_clim))

        # Adaptive thresholds
        limiar_sat = mediana_sat + k_sat * mad_sat
        limiar_clim = mediana_clim + k_clim * mad_clim

        # Default acceptance based on MAD thresholds
        mascara_erro = (grupo["Erro_abs_sat"] <= limiar_sat) | (grupo["Erro_abs_clim"] <= limiar_clim)

        # Additional acceptance: both Observed and Estimated <= 5 mm
        mascara_baixa_precip = (grupo["Observado"] <= 5) & (grupo["Estimado"] <= 5)

        # Combine both conditions
        mascara_aceita = mascara_erro | mascara_baixa_precip

        aceitas.append(grupo[mascara_aceita])
        rejeitadas.append(grupo[~mascara_aceita])

    return pd.concat(aceitas), pd.concat(rejeitadas)

# Apply the quality control function
df_aceitas, df_rejeitadas = controle_qualidade_regional(df)

# Save the results
df_aceitas.to_csv(args.saida_aceitas, sep="\t", index=False)
df_rejeitadas.to_csv(args.saida_rejeitadas, sep="\t", index=False)

print("Process completed.")
print(f"Accepted stations: {len(df_aceitas)}")
print(f"Rejected stations: {len(df_rejeitadas)}")