Recipe: Deep Shake v1.3

created by Serca

Profile

• Name: Deep Shake v1.3
• ID: 109225
• Shared with: Public
• Parent: Deep Shake
• Children:
  • X_Deep Shake v1.3
• Created on: October 27, 2025 at 10:56 AM UTC
• Updated on: October 29, 2025 at 10:55 AM UTC
• Description:

  Tests every rotamer and fuses the most promising. To test sidechains on specific segments, select them and check 'Test Selected Segments Only'.

Code

--[[

    v1.3 - Now you can run testing for selected segments only. 
         - Select segments and check the Checkbox "Test Selected Segments Only"
    v1.2 - Improved speed for the puzzles with the locked sidechains
	v1.1 - Added intermediate report every 300 tested rotamers


A variant of the Sidechain Flipper recipe. Unlike the standard version, which checks ~30% of rotamers, 
this one tries all of them, prioritizing the most promising. It may be more efficient, 
and even faster if stopped early after 10–20% yield no points.

See the discussion in the comments on this page for details: https://fold.it/recipes/109049
]]--

local SETTINGS = {
    SHAKE_ITERATIONS = 1,
    WIGGLE_ITERATIONS = 15,
    MAX_ROTAMER_ATTEMPTS = 10,
    MIN_ROTAMER_ATTEMPTS = 3,
    SERVICE_SLOT = 6,
    SERVICE_SLOT_TEMP = 7,
    SERVICE_SLOT_BEST = 5,
    MAX_SCORE_DECREASE = 10,
    FUZE_THRESHOLD_PERCENT = 50, -- top FUZE_THRESHOLD percents of tested rotamers will be fuzed
    FUZE_THRESHOLD = -1, -- points difference between bestScore and shake/Wiggle score to trigger FastFuze
    NORMALIZE_SUBSCORES = true, -- Enable/disable normalization of subscores for prioritization
    NORMALIZE_TOTAL_SCORE = true, -- Enable/disable normalization of total score
    NORMALIZE_ROTAMER_COUNT = true, -- Enable/disable normalization of rotamer count
    REPORT_CORRELATIONS_EVERY = 300, -- How often to print correlation table (0 = only at completion)
}
reportLevel = 2

-- Amino acids. Hydrophobics first.
single_letter_codes = { "g","a","v","l","i","m","f","w","p","s","t","c","y","n","q","d","e","k","r","h"}
three_letter_codes  = { "Gly","Ala","Val","Leu","Ile","Met","Phe","Trp","Pro","Ser","Thr","Cys","Tyr","Asn","Gln","Asp","Glu","Lys","Arg","His"}


fuzeConfig = {
    "{0.05, -7, 0.05,2} {0.25, 3, 1.00,20} {0.05, 3, 0.25,7} {0.25, 3, 0.25,20} {1.00, 3, 1.00,20}", --best score
    --"{0.25, 1, 1.00,20} {0.05, 2, 0.05,7} {0.25, 3, 0.05,7} {1.00, -2, 1.00,20} {1.00, 3, 1.00,20}", --best score
    "{0.25, -7, 0.25,7} {1.00, 2, 1.00,20}",	--fastest
    "{0.05, -20, 0.05,7} {0.25, 2, 0.25,7} {1.00, 2, 1.00,20}",	--fastest
    "{0.25, -7, 0.05,7} {0.05, 3, 0.25,20} {1.00, 2, 1.00,20}",
    "{1.00, 3, 1.00,7} {0.05, 1, 0.25,2} {1.00, 3, 1.00,20}",
    "{0.05, -2, 0.05,7} {1.00, 3, 1.00,7} {0.05, 2, 0.05,7} {0.25, 1, 0.25,2} {1.00, 3, 1.00,20}"
    }	--fastest
    
-- Just some corrrelations from a random puzzle to prioritize more perspective rotamers. 
-- Feel free to test for better ratios

correlations = {
    Packing             = -0.508,
    Clashing            = 0.508,
    Reference           = 0.270,
    Hiding              = 0.112,
    Bonding             = -0.106,
    Sidechain           = 0.090,
    Backbone            = 0.061,
    Disulfides          = -0.030,
    Pairwise            = -0.017,
    Ideality            = -0.012,

    TotalScore = 0.501,     -- Weight for EnergyScore
    RotamerCount = 0.185  -- Weight for rotamer count per segment
}

-- correlations = {
--     Hiding              = 0.306,
--     Ideality            = -0.253,
--     Backbone            = -0.249,
--     Packing             = -0.228,
--     Pairwise            = -0.151,
--     Clashing            = 0.135,
--     Bonding             = -0.081,
--     Sidechain           = 0.075,
--     Reference           = 0.043,
--     TotalScore = 0.5,     
--     RotamerCount = 0.185  
-- }


bestScore = current.GetEnergyScore()
print ("Starting with",bestScore)

function ScoreReturn()
	x = current.GetEnergyScore()
	return x-x%0.0001
end 
function roundx(score)
    if score == 0 then
        return 0
    end
    -- Round to 3 digits
    local rounded = math.floor(score * 1000 + 0.5) / 1000
    if rounded == 0 then
        -- Cheking order of magnitude: for score = 0.00005 log10(0.00005) ≈ -4.301, floor gives -5
        local exponent = math.floor(math.log10(math.abs(score)))
        local factor = 10^(-exponent) ---- Round to the first non-zero number
        rounded = math.floor(score * factor + 0.5) / factor
    end
    return rounded
end

function get_aa3 ( i )
    k = structure.GetAminoAcid (  i )

    for  j = 1, 20 do
        if ( k == single_letter_codes [  j ] ) then
            return three_letter_codes [  j  ]
        end
    end

    return "Unk"
end

-- Additional Fuze function to parse the input string into a 2D array
function parseInput(input)
    local result = {}
    for quadruplet in input:gmatch("{([^}]+)}") do
        local values = {}
        for value in quadruplet:gmatch("[^,%s]+") do
            table.insert(values, tonumber(value))
        end
        table.insert(result, {
            clashImportance = values[1],
            shakeIter = values[2],
            clashImportance2= values[3],
            wiggleIter = values[4]
        })
    end
    return result
end

-- Function for converting a correlation table into a string.
function correlationsToString(correlations)
    -- Fixed order of keys to display
    local orderedKeys = {
        "Packing", "Clashing", "Reference", "Hiding", 
        "Bonding", "Sidechain", "Backbone", "Disulfides", 
        "Pairwise", "Ideality", "TotalScore", "RotamerCount"
    }
    
    local result = "{"
    for _, name in ipairs(orderedKeys) do
        if correlations[name] ~= nil then
            result = result .. string.format("%s = %.3f, ", name, correlations[name])
        end
    end
    
    -- Add any remaining keys that weren't in our ordered list
    for name, value in pairs(correlations) do
        local found = false
        for _, orderedName in ipairs(orderedKeys) do
            if name == orderedName then
                found = true
                break
            end
        end
        
        if not found then
            result = result .. string.format("%s = %.3f, ", name, value)
        end
    end
    
    result = result .. "}"
    return result
end


-- Function for parsing a string of correlations.
function parseCorrelations(correlationString)
    local parsedCorrelations = {}
    
    -- Remove the curly braces at the beginning and end of the string, if they exist.
    correlationString = correlationString:gsub("^%s*{", ""):gsub("}%s*$", "")
    
    -- We split the string into pairs of name=value.
    for pair in correlationString:gmatch("[^,]+") do
        -- Removing extra spaces
        pair = pair:gsub("^%s+", ""):gsub("%s+$", "")
        
        -- Extracting the name and value.
        local name, value = pair:match("([%w_]+)%s*=%s*([-%d%.]+)")
        if name and value then
            parsedCorrelations[name] = tonumber(value)
        end
    end
    
    -- "We are checking that we were able to parse something."
    local count = 0
    for _, _ in pairs(parsedCorrelations) do count = count + 1 end
    
    if count == 0 then
        print("Error when parsing correlation string. Default values are being used.")
        return nil
    end
    
    return parsedCorrelations
end

-- Helper function for Pearson correlation calculation
function pearson_correlation(x, y)
    local n = #x
    local sum_x, sum_y, sum_xy, sum_x2, sum_y2 = 0, 0, 0, 0, 0
    for i = 1, n do
        sum_x = sum_x + x[i]
        sum_y = sum_y + y[i]
        sum_xy = sum_xy + x[i]*y[i]
        sum_x2 = sum_x2 + x[i]^2
        sum_y2 = sum_y2 + y[i]^2
    end
    local numerator = n*sum_xy - sum_x*sum_y
    local denominator = math.sqrt((n*sum_x2 - sum_x^2)*(n*sum_y2 - sum_y^2))
    return denominator ~= 0 and numerator/denominator or 0
end


-- Universal function that runs the logic on parsed input
function Fuze(inputString)
    recentbest.Save()
        save.Quicksave(99) 

        local fuzeConfig = parseInput(inputString)
        local score = ScoreReturn()

        for idx, triplet in ipairs(fuzeConfig) do
            if idx>1 then band.DisableAll() end
            behavior.SetClashImportance(triplet.clashImportance)
            if triplet.shakeIter and triplet.shakeIter > 0 then
                structure.ShakeSidechainsAll(triplet.shakeIter)
            elseif triplet.shakeIter and triplet.shakeIter < 0 then
                structure.WiggleAll(-triplet.shakeIter, false, true)
            end
            behavior.SetClashImportance(triplet.clashImportance2 )
            if triplet.wiggleIter and triplet.wiggleIter > 0 then
                structure.WiggleAll(triplet.wiggleIter)
            elseif triplet.wiggleIter and triplet.wiggleIter < 0 then
                structure.WiggleAll(-triplet.wiggleIter, true, false)
            end

            if triplet.clashImportance2==1 then 
                if  ScoreReturn() > score  then
                    score =  ScoreReturn()
                    save.Quicksave(99)
                else
                    save.Quickload(99)
                end
            end

        end
    recentbest.Restore()	
    return score -- Return the score
end

local function getSidechainScore(segment)
    return current.GetSegmentEnergySubscore(segment, "Sidechain") or 0
end


-------------------------------------------------------------------------------------------------

-- Global tracking variables
local globalRotamerList = {}
local processedRotamers = {}
local initialScores = {}

-- Table for statistics.
local segmentStats = {}

--[[  Collects all rotamers from all segments with their scores ]]
function CollectAllRotamers(report)
    local totalSegments = structure.GetCount()
    globalRotamerList = {}
    initialScores = {}
    segmentStats = {}  -- "Resetting the statistics with each rebuild."
    
    -- "Saving the names of sub-scores."
    local scoreparts = puzzle.GetPuzzleSubscoreNames()
    
    -- Save initial state
    save.Quicksave(SETTINGS.SERVICE_SLOT_TEMP)

    print ("Collecting data on available rotamers")
    
    -- Collect all rotamers
    for segment = 1, totalSegments do
        -- Skip segments whose sidechain is locked, or not selected when Test Selected Only is ON
        local _, sideLocked = structure.IsLocked(segment)
        local onlySelected = SETTINGS.TEST_SELECTED_ONLY == true
        if sideLocked then
            if reportLevel > 1 or report then 
                print(string.format("%d /%d: Segment %d: Sidechain locked", segment, totalSegments, segment))
            end
        elseif onlySelected and not selection.IsSelected(segment) then
            if reportLevel > 1 or report then 
                print(string.format("%d /%d: Segment %d: Not selected", segment, totalSegments, segment))
            end
        else
        local rotCount = rotamer.GetCount(segment)
        initialScores[segment] = getSidechainScore(segment)
        
        -- Initializing statistics recording for the segment.
        segmentStats[segment] = {
            aa = get_aa3(segment),
            rotamers = rotCount,
            improvements = 0,
            diff_score = -math.huge,
            initial_subscores = {},
            final_subscores = {},
            all_rotamer_scores = {},  -- Stores all the ratings of the rotamers.
            all_rotamer_subscores = {}, -- Stores sub-scores for each rotamer.
            all_rotamer_diff_scores = {},
            initial_total_score = current.GetEnergyScore()
        }
        
        -- "Saving the initial subscores for the current state of the segment."
        for jj = 1, #scoreparts do
            local sub = current.GetSegmentEnergySubscore(segment, scoreparts[jj]) or 0
            segmentStats[segment].initial_subscores[scoreparts[jj]] = roundx(sub)
        end
        
        local segmentScores = {}
        initialScore = current.GetEnergyScore()
        
        -- Initialize the storage of sub-scores for all rotamers in this segment.
        segmentStats[segment].all_rotamer_subscores = {}
        
        -- Evaluate each rotamer
        for rotIdx = 1, rotCount do
            if rotCount ~= rotamer.GetCount(segment) then
                print ("Unpredicted behavior: rotamer count has changed for segment "..segment.." from "..rotCount.." to "..rotamer.GetCount(segment))
                break
            end
            rotamer.SetRotamer(segment, rotIdx)
            local score = current.GetEnergyScore()
            
            -- Collecting subscores for this rotamer.
            local subScores = {}
            for jj = 1, #scoreparts do
                local sub = current.GetSegmentEnergySubscore(segment, scoreparts[jj]) or 0
                subScores[scoreparts[jj]] = sub
            end
            
            -- Saving the sub-scores for this rotamer.
            segmentStats[segment].all_rotamer_scores[rotIdx] = score
            segmentStats[segment].all_rotamer_subscores[rotIdx] = subScores
            segmentStats[segment].all_rotamer_diff_scores[rotIdx] = -math.huge
            
            table.insert(globalRotamerList, {
                segment = segment,
                rotIdx = rotIdx,
                score = score,
                sidechainScore = getSidechainScore(segment),
                rotamerCount = rotCount,
                subScores = subScores  -- Sub-scores for weighted ranking
            })
            segmentScores[rotIdx] = score

            -- Rare case: improved global best just by setting a rotamer (no Wiggle)
            if score > (bestScore or -math.huge) then 
                bestScore = score
                print("Improved to", ScoreReturn(), "at segment", segment, "rotamer", rotIdx)
                save.Quicksave(SETTINGS.SERVICE_SLOT_BEST)
                save.Quicksave(SETTINGS.SERVICE_SLOT_TEMP)
            end
        end
        
        local bestSegmentScore = math.max(unpack(segmentScores))
        if reportLevel > 1 or report then 
            print(string.format("%d /%d: Segment %d: Collected %d rotamers | Best score: %.3f", segment, totalSegments, segment, rotCount, bestSegmentScore))
        end
        
        -- Restore original state
        save.Quickload(SETTINGS.SERVICE_SLOT_TEMP)
        end -- if not sideLocked / not selected
    end
   
    -- Sort global list by weighted score
    sortRotamersByWeightedScore(globalRotamerList, correlations, 2)
end

--[[ Reports statistics about collected rotamers ]]
function ReportRotamerStats()
    print("\n=== Rotamer Collection Report ===")
    print("Total rotamers collected:", #globalRotamerList)
    
    local best = globalRotamerList[1]
    if best then 
        print(string.format("Best rotamer: Seg %d (%s) Rot %d | Score: %.3f", 
            best.segment, get_aa3(best.segment), best.rotIdx, best.score))
    end
    print("===============================\n")
end

-- Feature for updating the list of rotameters.
function UpdateRotamerList()
    local prevScores = {}
    for seg, score in pairs(initialScores) do
        prevScores[seg] = score
    end
    print("Rebuilding rotamer list...")
    CollectAllRotamers(false)  -- Reassembling with a new conformation
    
    -- Updating the processing based on previous assessments.
    local rotamersImproved = 0
    for j, newRot in ipairs(globalRotamerList) do
        if newRot.score > (prevScores[newRot.segment] or math.huge) then
            processedRotamers[j] = nil
            rotamersImproved = rotamersImproved + 1
        end
    end
    print("Number of rotamers to reevaluate:", rotamersImproved)
end

-- Change in the weighted grade calculation function.
function calculateWeightedScore(rotamer, correlations, subScoreStats)
    local weightedScore = 0
    
    -- Processing the general account.
    if rotamer.score and correlations.TotalScore then
        local value = rotamer.score
        -- Normalization of TotalScore, if enabled and there is statistics.
        if SETTINGS.NORMALIZE_TOTAL_SCORE and subScoreStats and subScoreStats["TotalScore"] then
            local stats = subScoreStats["TotalScore"]
            value = (value - stats.min) / stats.range
        end
        weightedScore = weightedScore + (value * correlations.TotalScore)  
    end
    
    -- Processing the number of rotameters.
    if rotamer.rotamerCount and correlations.RotamerCount then
        local value = rotamer.rotamerCount
        -- Normalization of RotamerCount, if enabled and statistics are available.
        if SETTINGS.NORMALIZE_ROTAMER_COUNT and subScoreStats and subScoreStats["RotamerCount"] then
            local stats = subScoreStats["RotamerCount"]
            value = (value - stats.min) / stats.range
        end
        weightedScore = weightedScore + (value * correlations.RotamerCount)  
    end
    
    -- Processing the other subscores
    for subScoreName, correlation in pairs(correlations) do
        -- We skip TotalScore and RotamerCount, as they have already been accounted for separately.
        if subScoreName ~= "TotalScore" and subScoreName ~= "RotamerCount" then
            if rotamer.subScores and rotamer.subScores[subScoreName] then
                local value = rotamer.subScores[subScoreName]
                
                -- Normalization, if enabled and there is statistics.
                if SETTINGS.NORMALIZE_SUBSCORES and subScoreStats and subScoreStats[subScoreName] then
                    local stats = subScoreStats[subScoreName]
                    value = (value - stats.min) / stats.range
                end
                
                -- Adding the contribution of the subscore, multiplied by its correlation.
                weightedScore = weightedScore + (value * correlation)
            end
        end
    end
    
    return weightedScore
end


-- Updating the sorting function to use statistics
function sortRotamersByWeightedScore(rotamerList, correlations, reportLevel)
    -- Collecting statistics for subscores if normalization is enabled
    local shouldCollectStats = SETTINGS.NORMALIZE_SUBSCORES or 
                              SETTINGS.NORMALIZE_TOTAL_SCORE or 
                              SETTINGS.NORMALIZE_ROTAMER_COUNT
    local subScoreStats = shouldCollectStats and collectSubScoreStats(rotamerList) or nil
    
    -- Outputting statistics for debugging
    if reportLevel >= 2 and subScoreStats then
        print("\n=== Normalization Statistics ===")
        -- We are displaying statistics for TotalScore and RotamerCount separately.
        for _, specialName in ipairs({"TotalScore", "RotamerCount"}) do
            local stats = subScoreStats[specialName]
            if stats then
                local minDisplay = (math.abs(stats.min) >= 1000) and string.format("%.0f", stats.min) or string.format("%.3f", stats.min)
                local maxDisplay = (math.abs(stats.max) >= 1000) and string.format("%.0f", stats.max) or string.format("%.3f", stats.max)
                local meanDisplay = (math.abs(stats.mean) >= 1000) and string.format("%.0f", stats.mean) or string.format("%.3f", stats.mean)
                local rangeDisplay = (math.abs(stats.range) >= 1000) and string.format("%.0f", stats.range) or string.format("%.3f", stats.range)
                print(string.format("%-20s: Min=%s, Max=%s, Mean=%s, Range=%s", 
                    specialName, minDisplay, maxDisplay, meanDisplay, rangeDisplay))
            end
        end
        
        -- We are displaying statistics for the other sub-scores.
        if SETTINGS.NORMALIZE_SUBSCORES then
            print("\n--- Subscore Statistics ---")
            for name, stats in pairs(subScoreStats) do
                if name ~= "TotalScore" and name ~= "RotamerCount" then
                    local minDisplay = (math.abs(stats.min) >= 1000) and string.format("%.0f", stats.min) or string.format("%.3f", stats.min)
                    local maxDisplay = (math.abs(stats.max) >= 1000) and string.format("%.0f", stats.max) or string.format("%.3f", stats.max)
                    local meanDisplay = (math.abs(stats.mean) >= 1000) and string.format("%.0f", stats.mean) or string.format("%.3f", stats.mean)
                    local rangeDisplay = (math.abs(stats.range) >= 1000) and string.format("%.0f", stats.range) or string.format("%.3f", stats.range)
                    print(string.format("%-20s: Min=%s, Max=%s, Mean=%s, Range=%s", 
                        name, minDisplay, maxDisplay, meanDisplay, rangeDisplay))
                end
            end
        end
        print("========================================\n")
    end
    
    -- Sort and print the information about the weights of the subscores correlations.
    if reportLevel >= 1 then
        local sortedCorrelations = {}
        for name, weight in pairs(correlations) do
            if weight > 0 then
                table.insert(sortedCorrelations, {name = name, weight = weight})
            end
        end
        table.sort(sortedCorrelations, function(a, b) return a.weight > b.weight end)
        
        print("Weight coefficients:")
        for _, entry in ipairs(sortedCorrelations) do
            print(string.format("%s: %.3f", entry.name, entry.weight))
        end
    end
    
    -- Calculating weighted score for each rotamer and saving it
    for _, rotamer in ipairs(rotamerList) do
        rotamer.weightedScore = calculateWeightedScore(rotamer, correlations, subScoreStats)
    end
    
    -- Sorting rotamers by weighted score (from highest to lowest)
    table.sort(rotamerList, function(a, b) 
        return a.weightedScore > b.weightedScore 
    end)
    
    return rotamerList
end

-- Add this feature to collect statistics on the sub-scores of all rotameters.
function collectSubScoreStats(rotamerList)
    local stats = {}
    
    -- Initialization of statistics for each type of subscore.
    for _, rotamer in ipairs(rotamerList) do
        if rotamer.subScores then
            for name, _ in pairs(rotamer.subScores) do
                if not stats[name] then
                    stats[name] = {min = math.huge, max = -math.huge, sum = 0, count = 0}
                end
            end
        end
    end
    
    -- Adding statistics for TotalScore and RotamerCount.
    stats["TotalScore"] = {min = math.huge, max = -math.huge, sum = 0, count = 0}
    stats["RotamerCount"] = {min = math.huge, max = -math.huge, sum = 0, count = 0}
    
    -- Collection of statistics
    for _, rotamer in ipairs(rotamerList) do
        -- "We are collecting statistics for subscores."
        if rotamer.subScores then
            for name, value in pairs(rotamer.subScores) do
                stats[name].min = math.min(stats[name].min, value)
                stats[name].max = math.max(stats[name].max, value)
                stats[name].sum = stats[name].sum + value
                stats[name].count = stats[name].count + 1
            end
        end
        
        -- "We are collecting statistics for TotalScore."
        if rotamer.score then
            stats["TotalScore"].min = math.min(stats["TotalScore"].min, rotamer.score)
            stats["TotalScore"].max = math.max(stats["TotalScore"].max, rotamer.score)
            stats["TotalScore"].sum = stats["TotalScore"].sum + rotamer.score
            stats["TotalScore"].count = stats["TotalScore"].count + 1
        end
        
        -- "We are collecting statistics for RotamerCount."
        if rotamer.rotamerCount then
            stats["RotamerCount"].min = math.min(stats["RotamerCount"].min, rotamer.rotamerCount)
            stats["RotamerCount"].max = math.max(stats["RotamerCount"].max, rotamer.rotamerCount)
            stats["RotamerCount"].sum = stats["RotamerCount"].sum + rotamer.rotamerCount
            stats["RotamerCount"].count = stats["RotamerCount"].count + 1
        end
    end
    
    -- Calculation of averages and range for all parameters.
    for name, data in pairs(stats) do
        data.mean = data.count > 0 and data.sum / data.count or 0
        -- For sub-scores with a very small range, we use a unit range.
        data.range = data.max - data.min
        if data.range < 0.0001 then
            data.range = 1
        end
    end
    
    return stats
end

--Major function with the Sidechain Flipper logic.
-----------------------------------------------------------------------------------------------
--[[ Modified main processing function ]]
function ProcessGlobalRotamers()
    local improvements = 0
    processedRotamers = {}
    diff_scores = {}
    
    -- Data points collection for correlation analysis
    local data_points = {}
    
    -- Adding a counter for unprocessed rotameters
    local unprocessedCount = #globalRotamerList
    local processedCount = 0
    
    -- Initial correlations - using what was provided by user
    local currentCorrelations = correlations
    
    while unprocessedCount > 0 do
        -- Find the first unprocessed rotamer
        local i = 1
        while i <= #globalRotamerList and processedRotamers[i] do
            i = i + 1
        end
        
        local rotData = globalRotamerList[i]
        -- Skip segments with 1 or fewer rotamers (locked sidechains were filtered during collection)
        if rotData.rotamerCount <= 1 then
            processedRotamers[i] = true
            unprocessedCount = unprocessedCount - 1
        else
            -- Set Rotamer to a new pozition, shake and wiggle a bit
            -- If improve is achieved, save the unfuzed version with the new rotamer
            save.Quicksave(SETTINGS.SERVICE_SLOT)
            local scoreBeforeRotamer = current.GetEnergyScore()
            if reportLevel > 2 then print ("Setting rotamer", rotData.segment, rotData.rotIdx) end

            freeze.Unfreeze(rotData.segment, false, true) --in case it is frozen
            rotamer.SetRotamer(rotData.segment, rotData.rotIdx)
            if current.GetEnergyScore() == scoreBeforeRotamer then  --if rotamer is the same as initial one, then do not process it
                save.Quickload(SETTINGS.SERVICE_SLOT)
            else
                freeze.Freeze(rotData.segment, false, true)
                
                structure.ShakeSidechainsAll(SETTINGS.SHAKE_ITERATIONS)
                --save unfuzed frozen rotamers version. to be restored if the score is improved
                save.Quicksave(SETTINGS.SERVICE_SLOT_TEMP) 
                freeze.UnfreezeAll()
                structure.WiggleAll(SETTINGS.WIGGLE_ITERATIONS)
                structure.ShakeSidechainsAll(SETTINGS.SHAKE_ITERATIONS)
                
                local diff = current.GetEnergyScore() - bestScore
                -- Start tracking the best delta achieved for this rotamer
                local finalDiff = diff
                
                if reportLevel > 2 then print(string.format("Score %s rotamer %d/%d diff %s", roundx(scoreBeforeRotamer),     rotData.rotIdx,  rotData.rotamerCount,    roundx(diff)))  end
                
                -- If score dropped by less than MAX_SCORE_DECREASE (i.e., within threshold)
                -- diff = current - best; drop < T  => diff > -T
                if diff > -SETTINGS.MAX_SCORE_DECREASE then
                    --structure.ShakeSidechainsAll(SETTINGS.SHAKE_ITERATIONS)
                    structure.WiggleAll(SETTINGS.WIGGLE_ITERATIONS)

                    -- Find how good the new diff_score is ranked between tested rotamers
                    local diff2 = current.GetEnergyScore() - bestScore
                    if diff2 > finalDiff then finalDiff = diff2 end
                    table.insert(diff_scores, diff2)
                    table.sort(diff_scores, function(a, b) return a > b end)
                    -- Find current rotamer rank
                    local current_rank = nil
                    for i, entry in ipairs(diff_scores) do
                        if entry == diff2 then
                            current_rank = i
                            break
                        end
                    end
                    
                    -- FastFuze only elements that are good enough: less than 1 point from bestScore or in top 50% delta score after shake/wiggle
                    local threshold = math.max(1, math.ceil(#diff_scores * SETTINGS.FUZE_THRESHOLD_PERCENT/100))
                    --print (#diff_scores, diff2, current_rank, threshold)
                    local temp_str = " "
                    if (diff2 > SETTINGS.FUZE_THRESHOLD) or (current_rank <= threshold) then
                        Fuze(fuzeConfig[2])
                        diff = ScoreReturn() - bestScore
                        temp_str = temp_str .."."
                    end
                    --temp_str = " ".current_rank .. "/".. threshold ..temp_str

                    -- Data point and stats will be saved after all stages are considered

                    if reportLevel > 1 then 
                        print("#" .. i .. "/" .. #globalRotamerList .. ": Seg " .. rotData.segment .. " (" .. get_aa3(rotData.segment)..")",
                            "Rot " .. rotData.rotIdx , "| Rotamers: " .. rotData.rotamerCount , "| Rank/Threshold: " .. current_rank.."/"..threshold ,
                            "| Start: " .. roundx(rotData.score) , " | Δ: " .. roundx(diff2)..temp_str)
                    end

                    local diff3 = current.GetEnergyScore() - bestScore
                    if diff3 > finalDiff then finalDiff = diff3 end
                    if diff3 > 0 then
                        bestScore = current.GetEnergyScore()
                        segmentStats[rotData.segment].improvements = segmentStats[rotData.segment].improvements + 1
                        
                        -- Remember that this rotamer led to an improvement
                        segmentStats[rotData.segment].improved_rotamers = segmentStats[rotData.segment].improved_rotamers or {}
                        segmentStats[rotData.segment].improved_rotamers[rotData.rotIdx] = true
                        
                        print("Score improved to", ScoreReturn(),"at segment", rotData.segment, "rotamer", rotData.rotIdx,"| Δ: ", roundx(diff3))
                        improvements = improvements + 1
                        save.Quicksave(SETTINGS.SERVICE_SLOT_BEST)
                        save.Quickload(SETTINGS.SERVICE_SLOT_TEMP)
                        save.Quicksave(SETTINGS.SERVICE_SLOT)
                    end
                end

                -- Save finalDiff for this rotamer and add to data_points (unconditionally)
                segmentStats[rotData.segment].diff_score = finalDiff
                segmentStats[rotData.segment].all_rotamer_diff_scores[rotData.rotIdx] = finalDiff

                local entry = {
                    diff_score = finalDiff,
                    rotamers = rotData.rotamerCount,
                    aa = get_aa3(rotData.segment),
                    total_score = rotData.score,
                    improved = finalDiff > 0,
                    subscores = rotData.subScores or {}
                }
                table.insert(data_points, entry)
            end -- if current.GetEnergyScore() == scoreBeforeRotamer
                
            save.Quickload(SETTINGS.SERVICE_SLOT)
            processedRotamers[i] = true
            unprocessedCount = unprocessedCount - 1
            processedCount = processedCount + 1
            
            -- Periodic correlation reporting every X processed rotamers (0 = only at completion)
            if SETTINGS.REPORT_CORRELATIONS_EVERY > 0 and 
                processedCount % SETTINGS.REPORT_CORRELATIONS_EVERY == 0 then
                -- print("\n=== Correlations after " .. processedCount .. " rotamers ===")
                analyzeCorrelations(data_points, true)
            end
        end -- if rotamer.GetCount(rotData.segment) <= 1
    end -- while unprocessedCount > 0
    
    -- Use the final processed data to create a recommendation for next run
    if #data_points > 0 then
        local finalCorrelations = analyzeCorrelations(data_points, false)
        print("\nRecommended correlations for next run:")
        print(correlationsToString(finalCorrelations))
    end
    
    return improvements
end

---------------------------------------------------------------------------------------------
-- Analyze stats after all the processing is finished
function AnalyzeSegmentStatistics()

    local scoreparts = puzzle.GetPuzzleSubscoreNames()
    local data_points = {}
    local subscore_sums = {}
    local rotamer_counts = {}
    local aa_counts = {}
    
    -- Collecting data on all rotameters from all segments
    for seg, stats in pairs(segmentStats) do
        -- Skip segments with only one rotamer (glycine)
        if stats.rotamers > 1 and stats.all_rotamer_subscores then
            for rotIdx, subscores in pairs(stats.all_rotamer_subscores) do
                local score = stats.all_rotamer_scores[rotIdx] or 0
                local diff_score = stats.all_rotamer_diff_scores[rotIdx] or -math.huge
                
                -- Skip rotamers with uninitialized diff_score (i.e., with unprocessed rotamers)
                if diff_score > -math.huge then
                    local entry = {
                        diff_score = diff_score,
                        rotamers = stats.rotamers,
                        aa = stats.aa,
                        total_score = score,
                        improved = (stats.improved_rotamers and stats.improved_rotamers[rotIdx]) or false,
                        subscores = {}
                    }
                    
                    -- Copy subscores for this rotamer
                    for name, val in pairs(subscores) do
                        entry.subscores[name] = val
                        subscore_sums[name] = subscore_sums[name] or {sum = 0, count = 0}
                        subscore_sums[name].sum = subscore_sums[name].sum + val
                        subscore_sums[name].count = subscore_sums[name].count + 1
                    end
                    
                    table.insert(data_points, entry)
                    rotamer_counts[stats.rotamers] = (rotamer_counts[stats.rotamers] or 0) + 1
                    aa_counts[stats.aa] = (aa_counts[stats.aa] or 0) + 1
                end
            end
        end
    end

    print("\n=== Segment Statistics Analysis ===")
    
    -- Get correlations from common function with verbose output
    local newCorrelations = analyzeCorrelations(data_points, true)

    -- AA success distribution analysis
    local aa_data = {}
    local aa_total = {}

    for seg, stats in pairs(segmentStats) do
        aa_total[stats.aa] = (aa_total[stats.aa] or 0) + 1
    end
    
    -- Collecting data on improvements
    for seg, stats in pairs(segmentStats) do
        if stats.diff_score > -math.huge then
            local aa = stats.aa
            aa_data[aa] = aa_data[aa] or {total = aa_total[aa], improvements = 0}
            aa_data[aa].improvements = aa_data[aa].improvements + stats.improvements
        end
    end

    -- Convert to a table and sort by the ratio of improvements
    local sorted_aa = {}
    for aa, data in pairs(aa_data) do
        local ratio = data.total > 0 and data.improvements/data.total or 0
        table.insert(sorted_aa, {
            aa = aa,
            total = data.total,
            improvements = data.improvements,
            ratio = ratio
        })
    end
    
    table.sort(sorted_aa, function(a,b) 
        return a.ratio > b.ratio  -- Sorting in descending order by ratio
    end)

    print("\nAA success distribution (sorted by improvement ratio):")
    print("AA  | Total | Improv | Ratio")
    print("-----------------------------")
    for _, entry in ipairs(sorted_aa) do
        print(string.format("%-3s | %5d | %6d | %.3f",
            entry.aa,
            entry.total,
            entry.improvements,
            entry.ratio))
    end

    -- Analysis of average subscores for successful/unsuccessful
    print("\nSubscore averages for successful segments:")
    local threshold = 0 -- The threshold for success in points earned after processing
    local successful = {}
    local unsuccessful = {}
    
    -- Initialize structures for each subscore
    for _, name in ipairs(scoreparts) do
        successful[name] = {sum = 0, count = 0}
        unsuccessful[name] = {sum = 0, count = 0}
    end

    -- Count successful and unsuccessful rotameters
    local successful_count = 0
    for _, dp in ipairs(data_points) do
        if dp.diff_score > threshold then
            successful_count = successful_count + 1
            for name, val in pairs(dp.subscores) do
                successful[name] = successful[name] or {sum = 0, count = 0}
                successful[name].sum = successful[name].sum + val
                successful[name].count = successful[name].count + 1
            end
        else
            for name, val in pairs(dp.subscores) do
                unsuccessful[name] = unsuccessful[name] or {sum = 0, count = 0}
                unsuccessful[name].sum = unsuccessful[name].sum + val
                unsuccessful[name].count = unsuccessful[name].count + 1
            end
        end
    end
    
    print("Total successful rotamers:", successful_count)
        if successful_count > 0 then
        -- Print the average values for successful and unsuccessful cases
        for _, name in ipairs(scoreparts) do
            local s_avg = successful[name] and successful[name].count > 0 and successful[name].sum/successful[name].count or 0
            local u_avg = unsuccessful[name] and unsuccessful[name].count > 0 and unsuccessful[name].sum/unsuccessful[name].count or 0
            print(string.format("%-20s: Successful %6.3f vs Unsuccessful %6.3f (Δ %+6.3f)",
                name, s_avg, u_avg, s_avg - u_avg))
        end
    end

    -- Adding the calculation of the average number of rotameters
    local total_rotamers = 0
    local total_segments = 0
    for rot, count in pairs(rotamer_counts) do
        total_rotamers = total_rotamers + (rot * count)
        total_segments = total_segments + count
    end
    local mean_rotamers = total_segments > 0 and (total_rotamers / total_segments) or 0
    print("Total rotamer_counts " .. total_segments, "Mean " .. mean_rotamers)
    
    -- print("\nRecommendations:")
    -- -- Correction of the error: correlations[1] is a table, not a number.
    -- local str_temp = ""
    -- if newCorrelationsArray[1] and newCorrelationsArray[1].corr > 0 then str_temp = " high "
    -- else str_temp = " low " end
    -- print("- Prioritize segments with"..str_temp..(newCorrelationsArray[1] and newCorrelationsArray[1].name or "N/A").." scores")
    -- --print("- Focus on segments with "..(sorted_aa[1] and sorted_aa[1].aa or "N/A").." amino acids")
    -- if sorted_aa[1] and sorted_aa[1].aa then
    --     print(string.format("- Focus on segments with %s amino acids", sorted_aa[1].aa))
    -- end
    -- if mean_rotamers > 0 then
    --     print(string.format("- Consider rotamer count > %d as higher potential", math.floor(mean_rotamers + 0.5)))
    -- else
    --     print("- Rotamer count data unavailable")
    -- end
    print("============================================\n")
end

----------------------------------------------------------------------------------------------------------

-- Function for requesting rotamer options
function RequestRotamerOptions()
    local ask = dialog.CreateDialog("Sidechain Flipper Settings")
    ask.shakeIter = dialog.AddSlider("Shake Iterations", SETTINGS.SHAKE_ITERATIONS, 0, 10, 0)
    ask.wiggleIter = dialog.AddSlider("Wiggle Iterations", SETTINGS.WIGGLE_ITERATIONS, 5, 30, 0)

    ask.maxScoreDecrease = dialog.AddSlider("Max Score Decrease", SETTINGS.MAX_SCORE_DECREASE, 0, 100, 1)
    ask.fuzeThresholdPct = dialog.AddSlider("Fuze Threshold %", SETTINGS.FUZE_THRESHOLD_PERCENT, 1, 100, 0)
    ask.fuzeThreshold = dialog.AddSlider("Fuze Score Threshold", SETTINGS.FUZE_THRESHOLD, -5, 5, 1)
    

    -- Adding settings for normalization
    ask.normalizeSubscores = dialog.AddCheckbox("Normalize Scores", SETTINGS.NORMALIZE_SUBSCORES)
    
    -- Adding settings for adaptive correlations
    -- Periodic correlation reporting
    ask.reportCorrEvery = dialog.AddSlider("Report correlations every", SETTINGS.REPORT_CORRELATIONS_EVERY, 0, 1000, 0)
    ask.reportLevel = dialog.AddSlider("ReportLevel", reportLevel, 1, 3, 0)
    
    -- Adding a field for entering correlations
    local correlationsString = correlationsToString(correlations)
    ask.correlationsInput = dialog.AddTextbox("Correlations", correlationsString)

    -- Limit processing to selected segments
    ask.testSelectedOnly = dialog.AddCheckbox("Test Selected Segments Only", SETTINGS.TEST_SELECTED_ONLY or false)

    ask.OK = dialog.AddButton("OK", 1)
    ask.Cancel = dialog.AddButton("Cancel", 0)

    if dialog.Show(ask) > 0 then
        SETTINGS.SHAKE_ITERATIONS = ask.shakeIter.value
        SETTINGS.WIGGLE_ITERATIONS = ask.wiggleIter.value
        SETTINGS.MAX_SCORE_DECREASE = ask.maxScoreDecrease.value
        SETTINGS.FUZE_THRESHOLD_PERCENT = ask.fuzeThresholdPct.value
        SETTINGS.FUZE_THRESHOLD = ask.fuzeThreshold.value
        reportLevel = ask.reportLevel.value
        
        -- Updating normalization settings
        SETTINGS.NORMALIZE_SUBSCORES = ask.normalizeSubscores.value
        
        -- Update selection-only mode
        SETTINGS.TEST_SELECTED_ONLY = ask.testSelectedOnly.value
        
        -- Updating adaptive correlation settings
        -- Update periodic correlation reporting
        SETTINGS.REPORT_CORRELATIONS_EVERY = ask.reportCorrEvery.value
        
        -- Parsing and applying correlations
        local parsedCorrelations = parseCorrelations(ask.correlationsInput.value)
        if parsedCorrelations then
            correlations = parsedCorrelations
            print("Correlations successfully updated.")
        end
        
        -- Display the selected settings for confirmation
        print("Settings updated:")
        print("Shake Iterations: " .. SETTINGS.SHAKE_ITERATIONS)
        print("Wiggle Iterations: " .. SETTINGS.WIGGLE_ITERATIONS)
        print("Fuze Threshold %: " .. SETTINGS.FUZE_THRESHOLD_PERCENT)
        print("Normalize Subscores: " .. tostring(SETTINGS.NORMALIZE_SUBSCORES))
        print("Normalize Total Score: " .. tostring(SETTINGS.NORMALIZE_TOTAL_SCORE))
        print("Normalize Rotamer Count: " .. tostring(SETTINGS.NORMALIZE_ROTAMER_COUNT))
        print("Report Correlations Every: " .. SETTINGS.REPORT_CORRELATIONS_EVERY)
        print("Test Selected Segments Only: " .. tostring(SETTINGS.TEST_SELECTED_ONLY))
        print("Total Score Weight: " .. (correlations.TotalScore or 0))
        print("Rotamer Count Weight: " .. (correlations.RotamerCount or 0))
        return true
    end
    return false
end

-- New function to analyze correlations based on the current data
function analyzeCorrelations(data_points, verbose)
    local newCorrelationsArray = {}
    
    -- Correlation of the number of rotamers with diff_score
    local x_rot, y_rot = {}, {}
    for _, dp in ipairs(data_points) do
        table.insert(x_rot, dp.rotamers)
        table.insert(y_rot, dp.diff_score)
    end
    local rot_corr = #data_points > 1 and pearson_correlation(x_rot, y_rot) or 0
    
    -- Correlation of the total score with diff_score
    local x_total, y_total = {}, {}
    for _, dp in ipairs(data_points) do
        table.insert(x_total, dp.total_score)
        table.insert(y_total, dp.diff_score)
    end
    local total_corr = #data_points > 1 and pearson_correlation(x_total, y_total) or 0
    
    -- Get the list of subscore names
    local scoreNames = {}
    if #data_points > 0 and data_points[1].subscores then
        for name, _ in pairs(data_points[1].subscores) do
            table.insert(scoreNames, name)
        end
    else
        scoreNames = puzzle.GetPuzzleSubscoreNames()
    end
    
    -- Correlation of subscores with diff_score
    for _, name in ipairs(scoreNames) do
        local x, y = {}, {}
        for _, dp in ipairs(data_points) do
            if dp.subscores and dp.subscores[name] ~= nil then
                table.insert(x, dp.subscores[name])
                table.insert(y, dp.diff_score)
            end
        end
        local corr = #x > 1 and pearson_correlation(x, y) or 0
        table.insert(newCorrelationsArray, {name = name, corr = corr})
    end
    
    -- Sorting by correlation strength
    table.sort(newCorrelationsArray, function(a,b) return math.abs(a.corr) > math.abs(b.corr) end)
    
    -- Convert to correlations format
    local newCorrelations = {}
    for i = 1, #newCorrelationsArray do
        if newCorrelationsArray[i] and newCorrelationsArray[i].name and newCorrelationsArray[i].corr then
            newCorrelations[newCorrelationsArray[i].name] = newCorrelationsArray[i].corr
        end
    end
    
    -- Adding TotalScore and RotamerCount
    newCorrelations["TotalScore"] = total_corr
    newCorrelations["RotamerCount"] = rot_corr
    
    -- Report if requested
    if verbose then
        -- print("\nSubscores correlations with the Δ points:")
        -- for i = 1, #newCorrelationsArray do
        --     print(string.format("%-20s = %.3f", newCorrelationsArray[i].name, newCorrelationsArray[i].corr))
        -- end
        
        print("New correlations:")
        print(correlationsToString(newCorrelations))
    end
    
    return newCorrelations
end

-- Reworked main function
function main()
    -- Request settings first; if cancelled, exit cleanly without starting
    local ok = RequestRotamerOptions()
    if not ok then
        print("Cancelled")
        return
    end
    
    save.Quicksave(1)
    behavior.SetClashImportance(1)
    
    freeze.UnfreezeAll()
    band.DisableAll()
    -- if no shake was done on start solution
    while true do 
        temp = ScoreReturn()
        structure.ShakeSidechainsAll(SETTINGS.SHAKE_ITERATIONS) 
        if temp <= ScoreReturn() then break end
    end
    save.Quicksave(1)

    save.Quicksave(SETTINGS.SERVICE_SLOT_BEST)
    save.Quicksave(SETTINGS.SERVICE_SLOT)
    save.Quicksave(SETTINGS.SERVICE_SLOT_TEMP)
    bestScore = current.GetEnergyScore()
    ScriptStartScore = current.GetEnergyScore()
    print("Score: " .. ScriptStartScore)
    print("\n=== Preprocessing: Collecting Rotamer Data ===")
    
    --recentbest.Save() -- sart saving best scores till the end of the script to restore them later

    CollectAllRotamers(true)
    ReportRotamerStats()

    --recentbest.Restore() -- restore best scores from the recentbest.save command

    print("\n=== Starting Rotamer Processing ===")

    improvements = ProcessGlobalRotamers()
    print("Total improvements:", improvements)
    
    save.Quickload(SETTINGS.SERVICE_SLOT_BEST)

    Cleanup("Finished")
end

cleanUpTrigger = false
function Cleanup(err)
    print(err)
    behavior.SetClashImportance(1)

    if cleanUpTrigger then -- Protection against cleanUp loop on crash
        print("Probably crash in AnalyzeSegmentStatistics(). Protecting from calling CleanUp again")
        return end  
    cleanUpTrigger = true --switching the trigger ON to prevent entering cleanUp again.

    -- if string.find(err, "Cancelled") then  -- Check if the error message contains "Cancelled"
    --     print("Ok, cancelled. Just wait for the last fuze.")
    -- end
    -- save.Quickload(SETTINGS.SERVICE_SLOT_BEST)
    -- prevScore = ScoreReturn()
    -- Fuze(fuzeConfig[1])
        
    -- if ScoreReturn() < prevScore then
    --     save.Quickload(SETTINGS.SERVICE_SLOT_BEST)
    -- end

    save.Quickload(1)
    save.Quicksave(SETTINGS.SERVICE_SLOT)
    save.Quickload(SETTINGS.SERVICE_SLOT_TEMP)
    save.Quickload(SETTINGS.SERVICE_SLOT_BEST)
    --recentbest.Restore() -- restore best scores from the recentbest.save command
    --print("Best solution restored. Current score:", ScoreReturn())
    AnalyzeSegmentStatistics()
end



xpcall ( main , Cleanup )