import matplotlib.pyplot as plt
import matplotlib.patches as patches
import os

# Dictionary with parameters for the Ecological Risk Assessment profile

ecotox_data = {
    # Algae
    "subcapitata_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    
    # Crustaceans
    "Crustaceans_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "magna_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    
    # Fish
    "Fish_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "Fathead_minnow_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "Bluegill_sunfish_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "Rainbow_trout_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "Sheepshead_minnow_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    
    # Tetrahymena pyriformis
    "pyriformis_toxicity_c": {"bounds": [0, 0.4, 0.7, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "pyriformis_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    
    # Bees and Birds
    "Honey_bee_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "Colinus_virginanus_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "Anas_platyrhynchos_toxicity": {"bounds": [0, 1, 3, 5], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    
    # Bioaccumulation and Degradation
    "BCF_c": {"bounds": [0, 0.4, 0.7, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "BCF": {"bounds": [0, 1, 2, 3], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "Biodegradability": {"bounds": [0, 0.4, 0.7, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]} 
}

def create_ecotox_bars_corrected():
    """
    Generates and saves the proportional bar charts for the Ecological Risk Assessment parameters.
    Handles multiline titles for long species names and increases bottom margin.
    """
    fig_width = 8 / 2.54
    fig_height = 1.3 / 2.54

    output_dir = "ADMET_Ecotox_Bars_Multiline"
    os.makedirs(output_dir, exist_ok=True)
    
    bg_color = '#FFFFFF'

    for param, info in ecotox_data.items():
        bounds = info["bounds"]
        colors = info["colors"]
        
        labels = [str(b) for b in bounds]

        fig, ax = plt.subplots(figsize=(fig_width, fig_height))
        fig.patch.set_facecolor(bg_color)
        ax.set_facecolor(bg_color)

        min_val = bounds[0]
        max_val = bounds[-1]

        for i in range(len(colors)):
            start = bounds[i]
            width = bounds[i+1] - bounds[i]
            rect = patches.Rectangle(
                (start, 0), width, 1, 
                facecolor=colors[i], edgecolor='none'
            )
            ax.add_patch(rect)

        ax.set_xlim(min_val, max_val)
        ax.set_ylim(0, 1)

        ax.set_xticks(bounds)
        ax.set_xticklabels(labels, rotation=45, ha='right', rotation_mode='anchor', fontsize=9)
        
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_visible(False)
        
        ax.spines['bottom'].set_position(('outward', 5))
        ax.spines['bottom'].set_linewidth(1)
        ax.tick_params(axis='x', direction='out', length=4, width=1, colors='black')
        ax.get_yaxis().set_visible(False)

        # Dynamic label formatting based on parameter type and text length
        raw_display = param.replace("_", " ")
        
        # SPECIFIC FIX FOR LONG ECOTOX NAMES: Break into 2 lines
        if "minnow toxicity" in raw_display:
            raw_display = raw_display.replace(" toxicity", "\ntoxicity")
        elif "trout toxicity" in raw_display:
            raw_display = raw_display.replace(" toxicity", "\ntoxicity")
        elif "sunfish toxicity" in raw_display:
            raw_display = raw_display.replace(" toxicity", "\ntoxicity")
        elif "virginanus toxicity" in raw_display:
            raw_display = raw_display.replace(" toxicity", "\ntoxicity")
        elif "platyrhynchos toxicity" in raw_display:
            raw_display = raw_display.replace(" toxicity", "\n toxicity")
        elif "bee toxicity" in raw_display:
             raw_display = raw_display.replace(" toxicity", "\ntoxicity")

        display_param = raw_display

        if param.endswith("_c") or param == "Biodegradability":
            display_param += " (Prob)"
        elif param == "BCF":
            display_param += " (log L/kg)"
        else:
            display_param += " (-log LC50/IGC50)"
            
        # Standardize standard font size to 11 for all short plots
        current_fontsize = 11
        bottom_margin = 0.45

        # Check for newline to adjust margin and size
        if "\n" in display_param:
            bottom_margin = 0.55 # Increase bottom margin to fit two lines
            current_fontsize = 9 # Force smaller size for multiline for better fit

        # Standard fallback logic for long non-multiline labels if any exist
        elif len(display_param) > 30:
            current_fontsize = 9
            
        ax.set_xlabel(display_param, fontsize=current_fontsize, labelpad=5, weight='bold')

        plt.subplots_adjust(bottom=bottom_margin)
        
        filename = os.path.join(output_dir, f"{param}_bar.png")
        plt.savefig(filename, dpi=300, bbox_inches='tight', facecolor=fig.get_facecolor())
        plt.close()

if __name__ == "__main__":
    print("Starting generation of Ecological Risk Assessment profile bars (Multiline Correction)...")
    create_ecotox_bars_corrected()
    print("Process completed successfully! Check the 'ADMET_Ecotox_Bars_Multiline' folder.")