from opcua import Server
from random import uniform, randint
from time import sleep
from datetime import datetime

def simulate_realistic(prev, target_range, max_change):
    target = uniform(*target_range)
    delta = target - prev
    change = max(min(delta, max_change), -max_change)
    return prev + change

server = Server()
server.set_endpoint("opc.tcp://0.0.0.0:4840/watertreatment/server/")
uri = "http://riverbend.opcsim"
idx = server.register_namespace(uri)

plant = server.nodes.objects.add_object(idx, "WaterTreatmentPlant")

# === Zone 1: Raw Water Intake ===
raw_flow = plant.add_variable(idx, "RawWaterFlowRate_LPM", 180.0)
intake_temp = plant.add_variable(idx, "IntakeWaterTemp_C", 14.0)
turbidity = plant.add_variable(idx, "Turbidity_NTU", 1.5)

# === Zone 2: Filtration ===
filter_in_pressure = plant.add_variable(idx, "FilterInPressure_kPa", 95.0)
filter_out_pressure = plant.add_variable(idx, "FilterOutPressure_kPa", 90.0)
filter_flow = plant.add_variable(idx, "FilteredFlow_LPM", 175.0)
backwash_valve = plant.add_variable(idx, "BackwashValveOpen", False)

# === Zone 3: Chemical Treatment ===
chlorine_ppm = plant.add_variable(idx, "ChlorineConcentration_mgL", 1.2)
ph_level = plant.add_variable(idx, "pH_Level", 7.4)
chemical_tank_level = plant.add_variable(idx, "ChemicalTankLevel_percent", 60.0)

# === Zone 4: Clean Water Storage ===
tank_level = plant.add_variable(idx, "CleanWaterTankLevel_percent", 75.0)
tank_temp = plant.add_variable(idx, "TankWaterTemp_C", 16.0)
tank_overflow = plant.add_variable(idx, "TankOverflowAlarm", False)

# === Zone 5: Distribution ===
pump1_status = plant.add_variable(idx, "Pump1Running", True)
pump2_status = plant.add_variable(idx, "Pump2Running", False)
zone1_pressure = plant.add_variable(idx, "Zone1Pressure_kPa", 185.0)
zone2_pressure = plant.add_variable(idx, "Zone2Pressure_kPa", 178.0)
generator_status = plant.add_variable(idx, "BackupGeneratorOnline", False)

# Make everything writable
for var in plant.get_variables():
    var.set_writable()

# Initial values
vals = {
    "raw_flow": 180.0, "intake_temp": 14.0, "turbidity": 1.5,
    "filter_in": 95.0, "filter_out": 90.0, "filter_flow": 175.0,
    "chlorine": 1.2, "ph": 7.4, "chem_level": 60.0,
    "tank_level": 75.0, "tank_temp": 16.0,
    "zone1_press": 185.0, "zone2_press": 178.0
}

server.start()
print(f"[{datetime.now()}] Full OPC UA Water Plant Simulator started.")

try:
    while True:
        # === Realistic updates ===
        vals["raw_flow"] = simulate_realistic(vals["raw_flow"], (160, 200), 1.5)
        vals["intake_temp"] = simulate_realistic(vals["intake_temp"], (12, 24), 1.2)
        vals["turbidity"] = simulate_realistic(vals["turbidity"], (0.8, 2.5), 0.1)

        vals["filter_in"] = simulate_realistic(vals["filter_in"], (90, 110), 1.0)
        vals["filter_out"] = simulate_realistic(vals["filter_out"], (80, 100), 1.0)
        vals["filter_flow"] = simulate_realistic(vals["filter_flow"], (150, 190), 1.2)

        vals["chlorine"] = simulate_realistic(vals["chlorine"], (0.5, 2.0), 0.1)
        vals["ph"] = simulate_realistic(vals["ph"], (5.5, 8.5), 0.1)
        vals["chem_level"] = simulate_realistic(vals["chem_level"], (30, 90), 1.0)

        vals["tank_level"] = simulate_realistic(vals["tank_level"], (10, 100), 8.0)
        vals["tank_temp"] = simulate_realistic(vals["tank_temp"], (12, 20), 0.5)

        vals["zone1_press"] = simulate_realistic(vals["zone1_press"], (160, 210), 2.0)
        vals["zone2_press"] = simulate_realistic(vals["zone2_press"], (155, 220), 2.0)

        # Set values
        raw_flow.set_value(round(vals["raw_flow"], 2))
        intake_temp.set_value(round(vals["intake_temp"], 2))
        turbidity.set_value(round(vals["turbidity"], 2))

        filter_in_pressure.set_value(round(vals["filter_in"], 2))
        filter_out_pressure.set_value(round(vals["filter_out"], 2))
        filter_flow.set_value(round(vals["filter_flow"], 2))

        chlorine_ppm.set_value(round(vals["chlorine"], 3))
        ph_level.set_value(round(vals["ph"], 2))
        chemical_tank_level.set_value(round(vals["chem_level"], 1))

        tank_level.set_value(round(vals["tank_level"], 1))
        tank_temp.set_value(round(vals["tank_temp"], 2))
        tank_overflow.set_value(vals["tank_level"] > 95)

        pump1_status.set_value(randint(0, 100) > 10)  # 90% chance running
        pump2_status.set_value(randint(0, 100) > 85)  # 15% chance running
        generator_status.set_value(randint(0, 200) > 198)  # Rare trigger
        zone1_pressure.set_value(round(vals["zone1_press"], 2))
        zone2_pressure.set_value(round(vals["zone2_press"], 2))

        #print(f"[{datetime.now()}] Values updated. Clean tank: {vals['tank_level']:.1f}%, Flow: {vals['raw_flow']:.1f} LPM")
        sleep(2)

except KeyboardInterrupt:
    print("Shutting down OPC UA Simulator...")
    server.stop()
