#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
extract_mrams_levels.py

Extrae datos de MRAMS desde archivos GrADS (.ctl + .dat).
Guarda un archivo de texto con 1 o varios niveles de la variable seleccionada.

Compatible con Python 3.x
"""

import struct

# =============================
# CONFIGURACIÓN DEL USUARIO
# =============================

ctl_file = "Jezero_Mars2020_Ls0-atm-S-g07.ctl"
dat_file = ctl_file.replace(".ctl", ".dat")

# Variable a extraer
varname = "press"   # presión
levels = [1]        # primer nivel termodinámico (14.5 m)

# Rango de tiempos (Fortran -> Python índices)
# 👇 ATENCIÓN: se ajusta para que salgan EXACTAMENTE 288 pasos
t_start = 290 - 1
t_end   = t_start + 288   # rango exclusivo → dará 288 valores

# Archivos de salida
output_floor = "MRAMS_Jezero_%s_14m_floor.txt" % varname
output_rim   = "MRAMS_Jezero_%s_14m_rim.txt"   % varname

# =============================
# PARSEAR EL CTL
# =============================

nx, ny, nz, nt = None, None, None, None
vars_info = []
inside_vars = False

with open(ctl_file) as f:
    for line in f:
        parts = line.strip().split()
        if not parts:
            continue
        key = parts[0].lower()
        if key == "xdef":
            nx = int(parts[1])
        elif key == "ydef":
            ny = int(parts[1])
        elif key == "zdef":
            nz = int(parts[1])
        elif key == "tdef":
            nt = int(parts[1])
        elif key == "vars":
            inside_vars = True
            continue
        elif key == "endvars":
            inside_vars = False
            continue
        elif inside_vars:
            vname = parts[0]
            vlevs = int(parts[1])
            vars_info.append((vname, vlevs))

# Buscar variable
offset_var = 0
target_var = None
for vname, vlevs in vars_info:
    if vname == varname:
        target_var = (offset_var, vlevs)
        break
    offset_var += vlevs

if target_var is None:
    raise Exception("Variable %s no encontrada en el ctl" % varname)

var_offset, var_levels = target_var

print("Archivo con dimensiones: nx=%d, ny=%d, nz=%d, nt=%d" % (nx, ny, nz, nt))
print("Variable %s encontrada con %d niveles" % (varname, var_levels))

# =============================
# FUNCIÓN DE EXTRACCIÓN
# =============================

def extract_point(ix, iy, outfile):
    """
    Extrae valores de MRAMS en (ix,iy) para los niveles indicados.
    """
    fdat = open(dat_file, "rb")
    with open(outfile, "w") as fout:
        fout.write("time_index " + " ".join(["%s_level%d" % (varname, l+1) for l in levels]) + "\n")

        for tt in range(t_start, t_end):  # rango exclusivo → 288 valores exactos
            row = [str(tt+1)]
            for lev in levels:
                pos = (tt * total_levels * field_size +
                       var_offset * field_size +
                       lev * field_size +
                       (iy * nx + ix))
                fdat.seek(pos * 4)  # float32 = 4 bytes
                val = struct.unpack("f", fdat.read(4))[0]
                row.append(str(val))
            fout.write(" ".join(row) + "\n")
    fdat.close()
    print("Archivo guardado:", outfile)

# =============================
# EXTRAER PARA DOS PUNTOS
# =============================

# Pedir coords al usuario
ix_floor = int(input("Índice X (floor): ")) - 1
iy_floor = int(input("Índice Y (floor): ")) - 1
ix_rim   = int(input("Índice X (rim): ")) - 1
iy_rim   = int(input("Índice Y (rim): ")) - 1

field_size = nx * ny
total_levels = sum(v[1] for v in vars_info)

# Extraer floor y rim
extract_point(ix_floor, iy_floor, output_floor)
extract_point(ix_rim, iy_rim, output_rim)