import time
from machine import I2C, Pin, ADC
from ssd1306 import SSD1306_I2C

# =========================
# SETUP
# =========================
i2c     = I2C(0, sda=Pin(4), scl=Pin(5), freq=400000)
adc_bat = ADC(Pin(26))
led     = Pin(25, Pin.OUT)

oled = SSD1306_I2C(128, 64, i2c)

# =========================

# =========================
BMP180_ADDR = 0x77

class BMP180:
    def __init__(self, i2c):
        self.i2c       = i2c
        self.B5        = 0
        self.sea_level = 101325
        self._load_calibration()

    def _rs(self, reg):
        d = self.i2c.readfrom_mem(BMP180_ADDR, reg, 2)
        v = (d[0] << 8) | d[1]
        return v - 65536 if v > 32767 else v

    def _ru(self, reg):
        d = self.i2c.readfrom_mem(BMP180_ADDR, reg, 2)
        return (d[0] << 8) | d[1]

    def _load_calibration(self):
        self.AC1 = self._rs(0xAA); self.AC2 = self._rs(0xAC)
        self.AC3 = self._rs(0xAE); self.AC4 = self._ru(0xB0)
        self.AC5 = self._ru(0xB2); self.AC6 = self._ru(0xB4)
        self.B1  = self._rs(0xB6); self.B2  = self._rs(0xB8)
        self.MC  = self._rs(0xBC); self.MD  = self._rs(0xBE)

    def _raw_temp(self):
        self.i2c.writeto_mem(BMP180_ADDR, 0xF4, b'\x2E')
        time.sleep_ms(5)
        d = self.i2c.readfrom_mem(BMP180_ADDR, 0xF6, 2)
        return (d[0] << 8) | d[1]

    def _raw_pressure(self):
        self.i2c.writeto_mem(BMP180_ADDR, 0xF4, b'\x34')
        time.sleep_ms(8)
        d = self.i2c.readfrom_mem(BMP180_ADDR, 0xF6, 3)
        return ((d[0] << 16) | (d[1] << 8) | d[2]) >> 8

    def _calc_B5(self):
        UT = self._raw_temp()
        X1 = ((UT - self.AC6) * self.AC5) >> 15
        X2 = (self.MC << 11) // (X1 + self.MD)
        self.B5 = X1 + X2

    def read_sea_level(self):
        
        self._calc_B5()
        UP = self._raw_pressure()
        B6 = self.B5 - 4000
        X1 = (self.B2 * ((B6 * B6) >> 12)) >> 11
        X2 = (self.AC2 * B6) >> 11
        X3 = X1 + X2
        B3 = (((self.AC1 * 4 + X3)) + 2) >> 2
        X1 = (self.AC3 * B6) >> 13
        X2 = (self.B1 * ((B6 * B6) >> 12)) >> 16
        X3 = ((X1 + X2) + 2) >> 2
        B4 = max(1, (self.AC4 * (X3 + 32768)) >> 15)
        B7 = (UP - B3) * 50000
        p  = (B7 * 2) // B4 if B7 < 0x80000000 else (B7 // B4) * 2
        return p

# =========================
# MPU6050
# =========================
MPU_ADDR = 0x68

def mpu_init():
    i2c.writeto_mem(MPU_ADDR, 0x6B, b'\x00')

def mpu_read():
    data = i2c.readfrom_mem(MPU_ADDR, 0x3B, 14)
    def s16(h, l):
        v = (h << 8) | l
        return v - 65536 if v > 32767 else v
    ax = round(s16(data[0], data[1]) / 16384.0, 2)
    ay = round(s16(data[2], data[3]) / 16384.0, 2)
    az = round(s16(data[4], data[5]) / 16384.0, 2)
    return ax, ay, az

# =========================
# ISA ATMOSPHERIC MODEL

# =========================
def sim_temperature(alt):
    """ ISA Standard Atmosphere"""
    if alt < 11000:
        return round(15.0 - 0.0065 * alt, 1)   # Troposphere
    elif alt < 20000:
        return -56.5                             # Tropopause
    else:
        return round(-56.5 + 0.001 * (alt - 20000), 1)  # Stratosphere

def sim_pressure(alt):
    """ ISA Standard Atmosphere"""
    if alt < 11000:
        p = 101325 * ((1 - 0.0000225577 * alt) ** 5.25588)
    elif alt < 20000:
        p = 22632 * (2.71828 ** (-0.0001577 * (alt - 11000)))
    else:
        p = 5474 * ((1 + 0.000004616 * (alt - 20000)) ** -34.163)
    return round(p)

def sim_altitude(pressure):
   
    return round(44330 * (1 - (pressure / 101325) ** (1 / 5.255)), 1)

# =========================
# GPS 
# =========================
MAX_ALT      = 30000
ASCENT_RATE  = 5
DESCENT_RATE = 8

def gps_balloon(t):
    ascent_time = MAX_ALT / ASCENT_RATE
    if t <= ascent_time:
        alt   = t * ASCENT_RATE
        phase = "ASCENT"
    else:
        descended = (t - ascent_time) * DESCENT_RATE
        alt       = max(0, MAX_ALT - descended)
        phase     = "DESCENT" if alt > 0 else "LANDED"
    lat = round(27.05 + t * 0.00003,  6)
    lon = round(30.05 + t * 0.000015, 6)
    return lat, lon, round(alt, 1), phase

# =========================
# BATTERY
# =========================
def read_battery():
    raw     = adc_bat.read_u16()
    voltage = round((raw / 65535) * 3.3, 2)
    pct     = max(0, min(100, int((voltage / 3.3) * 100)))
    return voltage, pct

# =========================
# CSV LOGGING
# =========================
LOG_FILE = "flight_log.csv"

def init_log():
    with open(LOG_FILE, "w") as f:
        f.write("time_s,temp_c,pressure_pa,altitude_m,lat,lon,gps_alt_m,phase,ax,ay,az,bat_v,bat_pct\n")

def log_data(t, temp, pressure, alt, lat, lon, gps_alt, phase, ax, ay, az, bat_v, bat_pct):
    with open(LOG_FILE, "a") as f:
        f.write(f"{t},{temp},{pressure},{alt},{lat},{lon},{gps_alt},{phase},{ax},{ay},{az},{bat_v},{bat_pct}\n")

# =========================
# OLED
# =========================
anim_frame = 0

def update_oled(temp, pressure, alt, lat, lon, bat_pct, phase, frame):
    oled.fill(0)

    dots = [".", "..", "..."]
    oled.text("HAB" + dots[frame % 3], 0, 0)
    oled.text(phase[:7], 72, 0)

    for x in range(128):
        oled.pixel(x, 9, 1)

    oled.text(f"T:{temp}C", 0, 12)
    oled.text(f"ALT:{int(alt)}m", 64, 12)
    oled.text(f"P:{pressure}Pa", 0, 23)
    oled.text(f"LA:{lat}", 0, 33)

    # Battery bar
    oled.text("BAT:", 0, 46)
    bar_w = int((bat_pct / 100) * 78)
    for x in range(78):
        oled.pixel(34 + x, 46, 1)
        oled.pixel(34 + x, 53, 1)
        if x < bar_w:
            for y in range(1, 7):
                oled.pixel(34 + x, 46 + y, 1)
    for y in range(2, 5):
        oled.pixel(113, 46 + y, 1)

    if frame % 2 == 0:
        oled.pixel(126, 0, 1); oled.pixel(127, 0, 1)
        oled.pixel(126, 1, 1); oled.pixel(127, 1, 1)

    oled.show()

# =========================
# BOOT ANIMATION
# =========================
def boot_animation():
    steps = ["BOOTING HAB...", "INIT SENSORS..", "CALIBRATING...", "SYSTEM READY! "]
    for i, msg in enumerate(steps):
        oled.fill(0)
        pct = int(((i + 1) / len(steps)) * 100)
        bar = int((pct / 100) * 100)
        for x in range(102):
            oled.pixel(13 + x, 38, 1)
            oled.pixel(13 + x, 45, 1)
        for y in range(7):
            oled.pixel(13,  38 + y, 1)
            oled.pixel(114, 38 + y, 1)
        for x in range(bar):
            for y in range(1, 6):
                oled.pixel(14 + x, 38 + y, 1)
        oled.text(msg, max(0, (128 - len(msg) * 8) // 2), 20)
        oled.text(f"{pct}%", 56, 50)
        oled.show()
        time.sleep_ms(700)

# =========================
# INIT
# =========================
boot_animation()
print("BOOTING HAB...")

bmp = BMP180(i2c)
mpu_init()
init_log()

time.sleep(1)
sea_level = bmp.read_sea_level()
print(f"Sea Level Pressure: {sea_level} Pa")
print("SYSTEM READY")
print("time_s | temp | pressure | alt | GPS | phase | accel | battery")

t0 = time.time()

# =========================
# MAIN LOOP
# =========================
while True:
    t = int(time.time() - t0)

    try:
        ax, ay, az = mpu_read()
    except Exception as e:
        print("Sensor error:", e)
        time.sleep(1)
        continue

    lat, lon, gps_alt, phase = gps_balloon(t)

   
    temp     = sim_temperature(gps_alt)
    pressure = sim_pressure(gps_alt)
    alt      = sim_altitude(pressure)

    bat_v, bat_pct = read_battery()

    # CSV log
    log_data(t, temp, pressure, alt, lat, lon, gps_alt,
             phase, ax, ay, az, bat_v, bat_pct)

    # Serial
    print(
        f"T:{t}s | {temp}C | {pressure}Pa | ALT:{alt}m | "
        f"GPS:{lat},{lon},{gps_alt} | {phase} | "
        f"AX:{ax},AY:{ay},AZ:{az} | BAT:{bat_v}V({bat_pct}%)"
    )

    # OLED
    update_oled(temp, pressure, alt, lat, lon, bat_pct, phase, anim_frame)
    anim_frame += 1

    led.toggle()
    time.sleep(2)