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

# Dictionary with continuous parameters for the Metabolism profile

metabolism_data = {
    # CYP Inhibitors (Traffic light colors: Green, Yellow, Red)
    "CYP1A2_inhibitor": {"bounds": [0, 0.2, 0.6, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "CYP3A4_inhibitor": {"bounds": [0, 0.15, 0.5, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "CYP2B6_inhibitor": {"bounds": [0, 0.25, 0.6, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "CYP2C9_inhibitor": {"bounds": [0, 0.15, 0.5, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "CYP2C19_inhibitor": {"bounds": [0, 0.15, 0.5, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "CYP2D6_inhibitor": {"bounds": [0, 0.15, 0.8, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    
    # CYP Substrates (Custom palette from the HTML table)
    "CYP1A2_substrate": {"bounds": [0, 0.2, 0.6, 1.0], "colors": ["#9FD7EF", "#DBDBDB", "#6787E7"]},
    "CYP3A4_substrate": {"bounds": [0, 0.3, 0.65, 1.0], "colors": ["#9FD7EF", "#DBDBDB", "#FFA333"]},
    "CYP2B6_substrate": {"bounds": [0, 0.2, 0.8, 1.0], "colors": ["#9FD7EF", "#DFCBDF", "#845CAD"]},
    "CYP2C9_substrate": {"bounds": [0, 0.2, 0.6, 1.0], "colors": ["#9FD7EF", "#7FB1B3", "#F39999"]},
    "CYP2C19_substrate": {"bounds": [0, 0.25, 0.75, 1.0], "colors": ["#9FD7EF", "#DBDBDB", "#EDAB94"]},
    "CYP2D6_substrate": {"bounds": [0, 0.2, 0.75, 1.0], "colors": ["#9FD7EF", "#DBDBDB", "#C777DB"]},
    
    # Liver Microsomes Stability (HLM / RLM)
    "HLM": {"bounds": [0, 0.3, 0.7, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "RLM": {"bounds": [0, 0.3, 0.7, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    
    # UGT Substrate (Inverted traffic light: Red, Yellow, Green)
    "UGT_substrate": {"bounds": [0, 0.3, 0.7, 1.0], "colors": ["#C14E4E", "#FFDF33", "#63C28D"]},
}

def create_metabolism_bars():
    """
    Generates and saves the proportional bar charts for the Metabolism parameters.
    """
    # Dimensions requested: 8 cm x 2 cm (converted to inches for matplotlib)
    fig_width = 8 / 2.54
    fig_height = 1.3 / 2.54

    # Create a new directory specific for Figure 3 (Metabolism)
    output_dir = "ADMET_Metabolism_Bars"
    os.makedirs(output_dir, exist_ok=True)
    
    # Background color matching the established paper style
    bg_color = '#ffffff'

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

        # Initialize figure
        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]

        # Draw each colored segment proportionally based on value distance
        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)

        # Axes scaling
        ax.set_xlim(min_val, max_val)
        ax.set_ylim(0, 1)

        # X-axis setup (ticks and labels)
        ax.set_xticks(bounds)
        ax.set_xticklabels(labels, rotation=45, ha='right', rotation_mode='anchor', fontsize=9)
        
        # Hide top, left and right spines (borders)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_visible(False)
        
        # Offset the bottom spine slightly downwards so it does not overlap the bar
        ax.spines['bottom'].set_position(('outward', 5))
        ax.spines['bottom'].set_linewidth(1)
        
        # Tick styling
        ax.tick_params(axis='x', direction='out', length=4, width=1, colors='black')
        
        # Completely hide the Y axis
        ax.get_yaxis().set_visible(False)

        # Parameter label below the X-axis
        display_param = param.replace("_", " ")
        ax.set_xlabel(display_param, fontsize=11, labelpad=5, weight='bold')

        # Adjust layout manually to prevent UserWarning about margins
        plt.subplots_adjust(bottom=0.45)
        
        # Save the figure
        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 Metabolism profile bars...")
    create_metabolism_bars()
    print("Process completed successfully! Check the 'ADMET_Metabolism_Bars' folder.")