from scriptengine import *

# Pump Station Generator Script for CODESYS with Auto Pump Selection
# This script generates a complete CODESYS project for a pump station with three pumps
# with automatic pump selection based on runtime

# Get the primary project or open a project
proj = projects.primary  # Use the currently open project

# Find the Application object
found = proj.find("Application", True)
app = found[0]

print("Starting Pump Station project generation with auto pump selection...")

# -------------------------------------------------------
# 1. Create Data Types (DUTs)
# -------------------------------------------------------

# Create a Structure DUT for pump data
pump_data_struct = app.create_dut("stPumpData", DutType.Structure)

# Define the structure
pump_data_struct.textual_declaration.replace("""TYPE stPumpData :
STRUCT
    // Status Inputs
    xEnable         : BOOL;    // Enable signal for the pump
    xManualMode     : BOOL;    // Manual operation mode
    xManualOperation: BOOL;    // Manual operation command (start button in manual mode)
    
    // Safety Inputs
    xPhsSeq         : BOOL;    // Phase sequence protection (TRUE = OK)
    xThermal        : BOOL;    // Thermal protection (TRUE = OK)
    xOverload       : BOOL;    // Overload protection (TRUE = OK)
    
    // Feedback Signals
    xContactor      : BOOL;    // Contactor feedback signal (TRUE = contactor activated)
    
    // Status Outputs
    xPumpOn         : BOOL;    // Pump control signal (TRUE = pump on)
    xFailure        : BOOL;    // Pump failure signal
    
    // Counters and Timers
    diRunTimeTotal  : DINT;    // Total runtime in seconds
    diStartCount    : DINT;    // Total number of starts
    rCurrentLevel   : REAL;    // Current liquid level
    rOnLevel        : REAL;    // Level to turn the pump on
    rOffLevel       : REAL;    // Level to turn the pump off
END_STRUCT
END_TYPE""")

print("Created stPumpData structure")

# Create an Enumeration DUT for pump status
pump_status_enum = app.create_dut("ePumpStatus", DutType.Enumeration)

# Define the enumeration
pump_status_enum.textual_declaration.replace("""TYPE ePumpStatus :
(
    DISABLED := 0,    // Pump is disabled
    READY := 1,       // Pump is ready to operate
    RUNNING := 2,     // Pump is running
    MANUAL := 3,      // Pump is in manual mode
    FAILURE := 4      // Pump has a failure
);
END_TYPE""")

print("Created ePumpStatus enumeration")

# -------------------------------------------------------
# 2. Create Global Variables Lists (GVLs)
# -------------------------------------------------------

# Create a Global Variable List for the pump station
gvl = app.create_gvl("GVL_PumpStation")

# Define the global variables
gvl.textual_declaration.replace("""VAR_GLOBAL
    // System Configuration
    g_xSystemEnabled    : BOOL := TRUE;    // System enable flag
    g_rLiquidLevel      : REAL;            // Current liquid level in meters
    
    // Pump Data Array (3 pumps)
    g_astPumpData       : ARRAY[1..3] OF stPumpData;  // Data for all pumps
    g_aePumpStatus      : ARRAY[1..3] OF ePumpStatus; // Status of all pumps
    
    // Configuration Parameters
    g_tConfirmationTime : TIME := T#3S;    // Confirmation time for contactor activation
    
    // Pump Level Configuration
    g_rPump1OnLevel     : REAL := 2.0;     // Level to turn pump 1 on
    g_rPump1OffLevel    : REAL := 1.0;     // Level to turn pump 1 off
    
    g_rPump2OnLevel     : REAL := 4.0;     // Level to turn pump 2 on
    g_rPump2OffLevel    : REAL := 3.0;     // Level to turn pump 2 off
    
    g_rPump3OnLevel     : REAL := 5.0;     // Level to turn pump 3 on
    g_rPump3OffLevel    : REAL := 4.0;     // Level to turn pump 3 off
    
    // System Status
    g_xHighLevelAlarm   : BOOL;            // High level alarm
    g_xLowLevelAlarm    : BOOL;            // Low level alarm
    g_xSystemFailure    : BOOL;            // System failure flag
    
    // Auto Pump Selection
    g_diMaxRuntimeDifference : DINT := 15; // Maximum runtime difference in seconds (1 hour)
    g_xAutoSelectPumps  : BOOL := TRUE;    // Enable automatic pump selection
    
    // Sorted pump indexes (by runtime)
    g_aiSortedPumps     : ARRAY[1..3] OF INT := [1, 2, 3]; // Default order
END_VAR""")

print("Created GVL_PumpStation global variables")

# Create a persistent Global Variable List for settings
persistent_gvl = app.create_persistentvars("GVL_Settings")

# Define the persistent global variables
persistent_gvl.textual_declaration.replace("""VAR_GLOBAL PERSISTENT
    // System Settings
    p_xAutoStartEnabled : BOOL := TRUE;    // Auto-start on power up
    p_tStartupDelay     : TIME := T#30S;   // Startup delay time
    
    // Alarm Thresholds
    p_rHighLevelAlarm   : REAL := 5.5;     // High level alarm threshold
    p_rLowLevelAlarm    : REAL := 0.5;     // Low level alarm threshold
    
    // Runtime Statistics
    p_adiTotalRunTime   : ARRAY[1..3] OF DINT;  // Total runtime for each pump in seconds
    p_adiTotalStarts    : ARRAY[1..3] OF DINT;  // Total starts for each pump
END_VAR""")

print("Created GVL_Settings persistent variables")

# -------------------------------------------------------
# 3. Create Function Block for Pump Control
# -------------------------------------------------------

# Create a Function Block for pump control
fb_pump_control = app.create_pou("FB_PumpControl")

# Define the declaration part of the Function Block
fb_pump_control.textual_declaration.replace("""FUNCTION_BLOCK FB_PumpControl
VAR_INPUT
    stPumpData      : stPumpData;   // Pump data structure
    rLiquidLevel    : REAL;         // Current liquid level
    tConfirmTime    : TIME;         // Confirmation time for contactor activation
    rOnLevel        : REAL;         // Level to turn the pump on
    rOffLevel       : REAL;         // Level to turn the pump off
END_VAR
VAR_OUTPUT
    xPumpOn         : BOOL;         // Pump control signal (TRUE = pump on)
    xFailure        : BOOL;         // Pump failure signal
    out_ePumpStatus : ePumpStatus;  // Current pump status
END_VAR
VAR
    fbConfirmTimer  : TON;          // Confirmation timer for contactor activation
    fbRunTimeTimer  : TON;          // Runtime timer (15 second precision)
    xSafetyOK       : BOOL;         // All safety inputs OK
    xRunTimeTick    : BOOL;         // 15-second tick for runtime counting
    tRunTimePrecision : TIME := T#15S; // Runtime counting precision
    xPrevPumpState  : BOOL;         // Previous pump state for edge detection
    xStartDetected  : BOOL;         // Start edge detected
END_VAR""")

# Define the implementation part of the Function Block
fb_pump_control.textual_implementation.replace("""// Check if all safety inputs are OK
xSafetyOK := stPumpData.xPhsSeq AND stPumpData.xThermal AND stPumpData.xOverload;

// Determine pump status based on inputs
IF NOT stPumpData.xEnable THEN
    out_ePumpStatus := ePumpStatus.DISABLED;
ELSIF xFailure THEN
    out_ePumpStatus := ePumpStatus.FAILURE;
ELSIF stPumpData.xManualMode THEN
    out_ePumpStatus := ePumpStatus.MANUAL;
ELSIF xPumpOn THEN
    out_ePumpStatus := ePumpStatus.RUNNING;
ELSE
    out_ePumpStatus := ePumpStatus.READY;
END_IF;

// Control logic with safety checks
IF stPumpData.xEnable AND xSafetyOK AND NOT xFailure THEN
    // Automatic mode: pump controlled by liquid level
    IF NOT stPumpData.xManualMode THEN
        IF (rLiquidLevel >= rOnLevel) THEN
            xPumpOn := TRUE;
        ELSIF (rLiquidLevel <= rOffLevel) THEN
            xPumpOn := FALSE;
        END_IF;
    ELSE
        // Manual mode: operator decides the pump start
        xPumpOn := stPumpData.xManualOperation;
    END_IF;
ELSE
    // Safety: turn off pump if disabled, safety issue, or failure
    xPumpOn := FALSE;
END_IF;

// Confirmation of pump activation via contactor
IF xPumpOn THEN
    fbConfirmTimer(IN := NOT stPumpData.xContactor, PT := tConfirmTime);
    
    // Check for confirmation failure
    IF fbConfirmTimer.Q THEN
        // No confirmation received within time: signal failure
        xFailure := TRUE;
        xPumpOn := FALSE;
    END_IF;
ELSE
    fbConfirmTimer(IN := FALSE);
    
    // Reset failure when pump is disabled and re-enabled
    IF xFailure AND NOT stPumpData.xEnable THEN
        xFailure := FALSE;
    END_IF;
END_IF;

// Detect pump start for counter increment
xStartDetected := xPumpOn AND NOT xPrevPumpState;
xPrevPumpState := xPumpOn;

// Runtime measurement (15-second precision)
fbRunTimeTimer(IN := xPumpOn, PT := tRunTimePrecision);
IF fbRunTimeTimer.Q THEN
    fbRunTimeTimer(IN := FALSE);
    xRunTimeTick := TRUE;
ELSE
    xRunTimeTick := FALSE;
END_IF;""")

print("Created FB_PumpControl function block")

# -------------------------------------------------------
# 4. Create Function Block for Pump Sorting
# -------------------------------------------------------

# Create a Function Block for pump sorting
fb_pump_sorter = app.create_pou("FB_PumpSorter")

# Define the declaration part of the Function Block
fb_pump_sorter.textual_declaration.replace("""FUNCTION_BLOCK FB_PumpSorter
VAR_INPUT
    adiPumpRuntimes : ARRAY[1..3] OF DINT;  // Runtime for each pump in seconds
    diMaxDifference : DINT;                 // Maximum allowed runtime difference
    xAutoSelect     : BOOL;                 // Enable automatic selection
END_VAR
VAR_OUTPUT
    aiSortedPumps   : ARRAY[1..3] OF INT;   // Sorted pump indexes (1-3)
END_VAR
VAR
    i, j            : INT;                  // Loop counters
    iTempIndex      : INT;                  // Temporary index for sorting
    diTempRuntime   : DINT;                 // Temporary runtime for sorting
    adiSortedRuntimes : ARRAY[1..3] OF DINT; // Sorted runtimes
    aiPumpIndexes   : ARRAY[1..3] OF INT;   // Pump indexes for sorting
END_VAR""")

# Define the implementation part of the Function Block
fb_pump_sorter.textual_implementation.replace("""// Initialize arrays
FOR i := 1 TO 3 DO
    adiSortedRuntimes[i] := adiPumpRuntimes[i];
    aiPumpIndexes[i] := i;
END_FOR;

// Only sort if auto-select is enabled
IF xAutoSelect THEN
    // Simple bubble sort to order pumps by runtime (ascending)
    FOR i := 1 TO 2 DO
        FOR j := 1 TO 3-i DO
            IF adiSortedRuntimes[j] > adiSortedRuntimes[j+1] THEN
                // Swap runtimes
                diTempRuntime := adiSortedRuntimes[j];
                adiSortedRuntimes[j] := adiSortedRuntimes[j+1];
                adiSortedRuntimes[j+1] := diTempRuntime;
                
                // Swap indexes
                iTempIndex := aiPumpIndexes[j];
                aiPumpIndexes[j] := aiPumpIndexes[j+1];
                aiPumpIndexes[j+1] := iTempIndex;
            END_IF;
        END_FOR;
    END_FOR;
    
    // Check if sorting is needed based on runtime difference
    IF (adiSortedRuntimes[3] - adiSortedRuntimes[1]) > diMaxDifference THEN
        // Use the sorted order
        aiSortedPumps := aiPumpIndexes;
    ELSE
        // Keep the current order if difference is not significant
        // But ensure we have valid values (1-3)
        FOR i := 1 TO 3 DO
            IF aiSortedPumps[i] < 1 OR aiSortedPumps[i] > 3 THEN
                aiSortedPumps[i] := i; // Default to sequential order
            END_IF;
        END_FOR;
    END_IF;
ELSE
    // If auto-select is disabled, use sequential order
    FOR i := 1 TO 3 DO
        aiSortedPumps[i] := i;
    END_FOR;
END_IF;""")

print("Created FB_PumpSorter function block")

# -------------------------------------------------------
# 5. Create Helper Function to Get Pump Position
# -------------------------------------------------------

# Create a function to find the position of a pump in the sorted array
get_pump_position = app.create_pou("GetPumpPosition")

# Define the declaration part of the function
get_pump_position.textual_declaration.replace("""FUNCTION GetPumpPosition : INT
VAR_INPUT
    aiSortedPumps : ARRAY[1..3] OF INT;  // Sorted pump indexes
    iPumpNumber   : INT;                 // Pump number to find (1-3)
END_VAR
VAR
    i : INT;                             // Loop counter
END_VAR""")

# Define the implementation part of the function
get_pump_position.textual_implementation.replace("""// Initialize return value to default position
GetPumpPosition := iPumpNumber;  // Default to same position if not found

// Search for the pump number in the sorted array
FOR i := 1 TO 3 DO
    IF aiSortedPumps[i] = iPumpNumber THEN
        GetPumpPosition := i;  // Return the position (1-3)
        RETURN;
    END_IF;
END_FOR;""")

print("Created GetPumpPosition function")

# -------------------------------------------------------
# 6. Create Main Program (PLC_PRG)
# -------------------------------------------------------

# Create the main program
plc_prg = app.create_pou("PLC_PRG", PouType.Program)

# Define the declaration part of the program
plc_prg.textual_declaration.replace("""PROGRAM PLC_PRG
VAR
    // Pump control function blocks
    fbPump1 : FB_PumpControl;
    fbPump2 : FB_PumpControl;
    fbPump3 : FB_PumpControl;
    
    // Pump sorter function block
    fbPumpSorter : FB_PumpSorter;
    
    // Level configuration for sorted pumps
    arSortedOnLevels  : ARRAY[1..3] OF REAL;
    arSortedOffLevels : ARRAY[1..3] OF REAL;
    
    // Pump positions in sorted array
    iPump1Position : INT;
    iPump2Position : INT;
    iPump3Position : INT;
    
    // Initialization flag
    xInitialized : BOOL := FALSE;
    
    // Runtime counters
    fbRuntimeCounter : TON;
    tRuntimeInterval : TIME := T#15S;
    
    // String variables for pump status display
    sPump1 : WSTRING;
    sPump2 : WSTRING;
    sPump3 : WSTRING;
    
    // Array of pump state names
    sPumpStates : ARRAY[0..4] OF STRING := 
        ['DISABLED', 'READY', 'RUNNING', 'MANUAL', 'FAILURE'];
    
    // String variables for each pump state
    sPumpState1 : STRING;
    sPumpState2 : STRING;
    sPumpState3 : STRING;
    
    // Temporary variables
    i : INT;
END_VAR""")

# Define the implementation part of the program
plc_prg.textual_implementation.replace("""// Initialize pump data on first scan
IF NOT xInitialized THEN
    // Initialize Pump 1
    g_astPumpData[1].xEnable := TRUE;
    g_astPumpData[1].diRunTimeTotal := p_adiTotalRunTime[1];
    g_astPumpData[1].diStartCount := p_adiTotalStarts[1];
    
    // Initialize Pump 2
    g_astPumpData[2].xEnable := TRUE;
    g_astPumpData[2].diRunTimeTotal := p_adiTotalRunTime[2];
    g_astPumpData[2].diStartCount := p_adiTotalStarts[2];
    
    // Initialize Pump 3
    g_astPumpData[3].xEnable := TRUE;
    g_astPumpData[3].diRunTimeTotal := p_adiTotalRunTime[3];
    g_astPumpData[3].diStartCount := p_adiTotalStarts[3];
    
    // Initialize sorted pump array with default values (1,2,3)
    g_aiSortedPumps[1] := 1;
    g_aiSortedPumps[2] := 2;
    g_aiSortedPumps[3] := 3;
    
    xInitialized := TRUE;
END_IF;

// Check for alarms
g_xHighLevelAlarm := g_rLiquidLevel >= p_rHighLevelAlarm;
g_xLowLevelAlarm := g_rLiquidLevel <= p_rLowLevelAlarm;

// Update current level for all pumps
g_astPumpData[1].rCurrentLevel := g_rLiquidLevel;
g_astPumpData[2].rCurrentLevel := g_rLiquidLevel;
g_astPumpData[3].rCurrentLevel := g_rLiquidLevel;

// Run pump sorter to determine the order of pumps based on runtime
fbPumpSorter(
    adiPumpRuntimes := p_adiTotalRunTime,
    diMaxDifference := g_diMaxRuntimeDifference,
    xAutoSelect := g_xAutoSelectPumps
);
g_aiSortedPumps := fbPumpSorter.aiSortedPumps;

// Find positions of each pump in the sorted array
iPump1Position := GetPumpPosition(g_aiSortedPumps, 1);
iPump2Position := GetPumpPosition(g_aiSortedPumps, 2);
iPump3Position := GetPumpPosition(g_aiSortedPumps, 3);

// Assign level thresholds based on sorted pump order
// Pump with least runtime (index 1) gets the lowest level
arSortedOnLevels[1] := g_rPump1OnLevel;   // First pump (lowest runtime) - ON at 2m
arSortedOffLevels[1] := g_rPump1OffLevel; // First pump (lowest runtime) - OFF at 1m

arSortedOnLevels[2] := g_rPump2OnLevel;   // Second pump - ON at 4m
arSortedOffLevels[2] := g_rPump2OffLevel; // Second pump - OFF at 3m

arSortedOnLevels[3] := g_rPump3OnLevel;   // Third pump (highest runtime) - ON at 5m
arSortedOffLevels[3] := g_rPump3OffLevel; // Third pump (highest runtime) - OFF at 4m

// Control Pump 1
fbPump1(
    stPumpData := g_astPumpData[1],
    rLiquidLevel := g_rLiquidLevel,
    tConfirmTime := g_tConfirmationTime,
    rOnLevel := arSortedOnLevels[iPump1Position],
    rOffLevel := arSortedOffLevels[iPump1Position]
);
g_astPumpData[1].xPumpOn := fbPump1.xPumpOn;
g_astPumpData[1].xFailure := fbPump1.xFailure;
g_aePumpStatus[1] := fbPump1.out_ePumpStatus;

// Convert pump status to string for display
sPumpState1 := sPumpStates[TO_INT(fbPump1.out_ePumpStatus)];

// Control Pump 2
fbPump2(
    stPumpData := g_astPumpData[2],
    rLiquidLevel := g_rLiquidLevel,
    tConfirmTime := g_tConfirmationTime,
    rOnLevel := arSortedOnLevels[iPump2Position],
    rOffLevel := arSortedOffLevels[iPump2Position]
);
g_astPumpData[2].xPumpOn := fbPump2.xPumpOn;
g_astPumpData[2].xFailure := fbPump2.xFailure;
g_aePumpStatus[2] := fbPump2.out_ePumpStatus;

// Convert pump status to string for display
sPumpState2 := sPumpStates[TO_INT(fbPump2.out_ePumpStatus)];

// Control Pump 3
fbPump3(
    stPumpData := g_astPumpData[3],
    rLiquidLevel := g_rLiquidLevel,
    tConfirmTime := g_tConfirmationTime,
    rOnLevel := arSortedOnLevels[iPump3Position],
    rOffLevel := arSortedOffLevels[iPump3Position]
);
g_astPumpData[3].xPumpOn := fbPump3.xPumpOn;
g_astPumpData[3].xFailure := fbPump3.xFailure;
g_aePumpStatus[3] := fbPump3.out_ePumpStatus;

// Convert pump status to string for display
sPumpState3 := sPumpStates[TO_INT(fbPump3.out_ePumpStatus)];

// Set current pump for display (sorted pumps)
sPump1 := INT_TO_WSTRING(g_aiSortedPumps[1]);
sPump2 := INT_TO_WSTRING(g_aiSortedPumps[2]);
sPump3 := INT_TO_WSTRING(g_aiSortedPumps[3]);

// Update runtime counters (15-second precision)
fbRuntimeCounter(IN := TRUE, PT := tRuntimeInterval);
IF fbRuntimeCounter.Q THEN
    fbRuntimeCounter(IN := FALSE);
    
    // Update runtime for each running pump
    IF g_astPumpData[1].xPumpOn THEN
        g_astPumpData[1].diRunTimeTotal := g_astPumpData[1].diRunTimeTotal + 15;
        p_adiTotalRunTime[1] := g_astPumpData[1].diRunTimeTotal;
    END_IF;
    
    IF g_astPumpData[2].xPumpOn THEN
        g_astPumpData[2].diRunTimeTotal := g_astPumpData[2].diRunTimeTotal + 15;
        p_adiTotalRunTime[2] := g_astPumpData[2].diRunTimeTotal;
    END_IF;
    
    IF g_astPumpData[3].xPumpOn THEN
        g_astPumpData[3].diRunTimeTotal := g_astPumpData[3].diRunTimeTotal + 15;
        p_adiTotalRunTime[3] := g_astPumpData[3].diRunTimeTotal;
    END_IF;
END_IF;

// Update start counters on rising edge of pump activation
IF fbPump1.xStartDetected THEN
    g_astPumpData[1].diStartCount := g_astPumpData[1].diStartCount + 1;
    p_adiTotalStarts[1] := g_astPumpData[1].diStartCount;
END_IF;

IF fbPump2.xStartDetected THEN
    g_astPumpData[2].diStartCount := g_astPumpData[2].diStartCount + 1;
    p_adiTotalStarts[2] := g_astPumpData[2].diStartCount;
END_IF;

IF fbPump3.xStartDetected THEN
    g_astPumpData[3].diStartCount := g_astPumpData[3].diStartCount + 1;
    p_adiTotalStarts[3] := g_astPumpData[3].diStartCount;
END_IF;

// Set system failure if any pump has failed
g_xSystemFailure := g_astPumpData[1].xFailure OR g_astPumpData[2].xFailure OR g_astPumpData[3].xFailure;""")

print("Created PLC_PRG main program")

# -------------------------------------------------------
# 7. Create Task Configuration
# -------------------------------------------------------

# Create task configuration
tc = app.create_task_configuration()

# Add main task
task = tc.create_task("MainTask")
task.interval = "T#100MS"
task.priority = "1"
task.pous.add("PLC_PRG")

print("Created task configuration")

# Add main task
task = tc.create_task("VISU_TASK")
task.interval = "T#100MS"
task.priority = "1"
task.pous.add("VisuElems.Visu_Prg")

print("Created visu task configuration")


print("Pump Station project generation with auto pump selection completed successfully!")
print("The following components were created:")
print("- Data Types: stPumpData, ePumpStatus")
print("- Global Variables: GVL_PumpStation, GVL_Settings")
print("- Function Blocks: FB_PumpControl, FB_PumpSorter")
print("- Functions: GetPumpPosition")
print("- Programs: PLC_PRG with automatic pump selection")
print("- Task Configuration: MainTask")