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

# Dictionary with parameters for the Distribution profile
distribution_data = {
    # Custom BBB (Prob) 
    "BBB": {"bounds": [0, 0.3, 0.8, 1.0], "colors": ["#C14E4E", "#FFDF33", "#63C28D"]},
    
    # Transporters
    "OATP1B1_inhibitor": {"bounds": [0, 0.5, 0.8, 1.0], "colors": ["#DBDBDB", "#FFB733", "#C14E4E"]},
    "OATP1B3_inhibitor": {"bounds": [0, 0.6, 0.85, 1.0], "colors": ["#DBDBDB", "#FF926A", "#C14E4E"]},
    "OATP2B1_inhibitor": {"bounds": [0, 0.25, 0.45, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "OCT1_inhibitor": {"bounds": [0, 0.25, 0.55, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "OCT2_inhibitor": {"bounds": [0, 0.15, 0.45, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "BCRP_inhibitor": {"bounds": [0, 0.25, 0.6, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "BSEP_inhibitor": {"bounds": [0, 0.25, 0.7, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "MATE1_inhibitor": {"bounds": [0, 0.55, 1.0], "colors": ["#63C28D", "#C14E4E"]}, # Only 2 segments
    "Pgp_inhibitor": {"bounds": [0, 0.2, 0.6, 1.0], "colors": ["#63C28D", "#FFDF33", "#C14E4E"]},
    "Pgp_substrate": {"bounds": [0, 0.3, 0.7, 1.0], "colors": ["#3CB371", "#FFD700", "#FF8C00"]},
    
    # PPB (Ratio)
    "PPB": {"bounds": [0, 0.2, 0.7, 1.0], "colors": ["#C14E4E", "#FFDF33", "#63C28D"]},
    
    # Continuous Parameter VDss (-log L/kg)
    "VDss": {"bounds": [-1.5, -0.6, 0.5, 1.5], "colors": ["#C14E4E", "#FFDF33", "#4A90E2"]}
}

def create_distribution_bars():
    """
    Generates and saves the proportional bar charts for ADMET Distribution parameters.
    Handles exact custom formatting for bounds, diverse color palettes, and specific units.
    """
    fig_width = 8 / 2.54
    fig_height = 1.3 / 2.54

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

    for param, info in distribution_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
        raw_display = param.replace("_", " ")
        
        if " inhibitor" in raw_display:
            raw_display = raw_display.replace(" inhibitor", "\ninhibitor")
        elif " substrate" in raw_display:
            raw_display = raw_display.replace(" substrate", "\nsubstrate")

        # Specific units
        if param == "VDss":
            display_param = f"{raw_display} (-log L/kg)"
        elif param == "PPB":
            display_param = f"{raw_display} (Ratio)"
        else:
            display_param = f"{raw_display} (Prob)"
            
        current_fontsize = 11
        bottom_margin = 0.45

        if "\n" in display_param:
            bottom_margin = 0.55
            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 EXACT Distribution profile bars...")
    create_distribution_bars()
    print("Process completed successfully! Check the 'ADMET_Distribution_Bars' folder.")