import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

# Load the coverage data for all 8 datasets
positive_strands = []
negative_strands = []
for i in range(1, 9):
    positive_strands.append(np.load(f'positive_strand_coverage{i}.npy'))
    negative_strands.append(np.load(f'negative_strand_coverage{i}.npy') * -1)  # Invert negative strand values

# Load filtered data
positive_filtered = np.load('positive_filtered.npy')
negative_filtered = np.load('negative_filtered.npy') * -1  # Invert for plotting

# Load slide detector data
slide_detector_pos = np.load('slide-detector-pos.npy')
slide_detector_neg = np.load('slide-detector-neg.npy') * -1  # Invert negative strand values

# Start-end coordinates
cus_start, cus_end = 1,10000

# ---------- Prep x-axis and subset ----------
x_start, x_end = cus_start, cus_end
x = np.arange(x_start, x_end)

# Subset the data for all datasets
pos_subsets = [pos[x_start:x_end] for pos in positive_strands]
neg_subsets = [neg[x_start:x_end] for neg in negative_strands]

# Subset filtered data
pos_filtered_subset = positive_filtered[x_start:x_end]
neg_filtered_subset = negative_filtered[x_start:x_end]

# Subset slide detector data
slide_detector_pos_subset = slide_detector_pos[x_start:x_end]
slide_detector_neg_subset = slide_detector_neg[x_start:x_end]

# Calculate global maximum for rectangle height
global_max_coverage = max(
    max([subset.max() for subset in pos_subsets]),
    max([abs(subset.min()) for subset in neg_subsets]),
    pos_filtered_subset.max(),
    abs(neg_filtered_subset.min())
) * 1.05

# ---------- Create the plot ----------
fig, ax1 = plt.subplots(figsize=(12, 6))

# Define dataset groups with colors and labels
datasets = [
    {'index': 0, 'group': 'IVT (1)', 'color': '#0000ff'},  # Dark blue
    {'index': 1, 'group': 'IVT (2)', 'color': '#7755ff'},  # Light blue
    {'index': 2, 'group': 'IVT 10x DNA (3)', 'color': '#66aaee'},  # Green-blue
    {'index': 3, 'group': 'No polymerase (4)', 'color': '#d62728'},  # Dark red
    {'index': 4, 'group': 'No polymerase (5)', 'color': '#ff7f0e'},  # Medium red
    {'index': 5, 'group': 'No polymerase 10x DNA (6)', 'color': '#ff0088'},  # Light red
    {'index': 6, 'group': 'DNA (7)', 'color': '#17fe1f'},  # Dark green
    {'index': 7, 'group': 'DNA (8)', 'color': '#9efa15'}   # Light green
]

# Plot all datasets
for dataset in datasets:
    i = dataset['index']
    ax1.plot(x, pos_subsets[i], color=dataset['color'], linewidth=1, label=dataset['group'])
    ax1.plot(x, neg_subsets[i], color=dataset['color'], linewidth=1)

regions = [
    #{'start': cus_start, 'end': cus_end, 'color': 'lightgreen',   'label': 'kgtP(neg)'},
    #{'start': 1302899, 'end': 1303819, 'color': 'lightblue',    'label': 'oppB'},
    #{'start': 1301182, 'end': 1306768, 'color': "khaki",        'label': 'oppA-oppF'},
    #{'start': 1303834, 'end': 1304742, 'color': 'orange',         'label': 'oppC'},
    #{'start': 1300856, 'end': 1300915, 'color': 'blue',    'label': 'oppAp2'},
    #{'start': 1301025, 'end': 1301084, 'color': 'plum',    'label': 'oppAp'},
]

for r in regions:
    # Positive strand rectangle
    ax1.add_patch(
        Rectangle(
            (r['start'], 0),
            r['end'] - r['start'],
            global_max_coverage,
            facecolor=r['color'],
            alpha=0.75,
            label=r['label']
        )
    )
    # Negative strand rectangle
    ax1.add_patch(
        Rectangle(
            (r['start'], -global_max_coverage),
            r['end'] - r['start'],
            global_max_coverage,
            facecolor=r['color'],
            alpha=0.75
        )
    )

# Plot filtered data with grey fill
ax1.plot(x, pos_filtered_subset, color='grey', linewidth=1, label='Filtered')
ax1.plot(x, neg_filtered_subset, color='grey', linewidth=1)
ax1.fill_between(x, pos_filtered_subset, neg_filtered_subset, color='grey', alpha=0.7)

# Set up the primary y-axis
ax1.set_ylim(bottom=-global_max_coverage, top=global_max_coverage)
ax1.axhline(0, color='black', linewidth=0.5)  # Add zero line
ax1.set_ylabel('Coverage (positive and negative strands)')

# Create secondary y-axis for slide detector data
ax2 = ax1.twinx()

# Calculate max for slide detector data for scaling
slide_detector_max = max(
    slide_detector_pos_subset.max(),
    abs(slide_detector_neg_subset.min())
) * 1.05

# Plot slide detector data on secondary axis
ax2.plot(x, slide_detector_pos_subset, color='purple', linewidth=1, label='Slide Detector (+)')
ax2.plot(x, slide_detector_neg_subset, color='magenta', linewidth=1, label='Slide Detector (-)')
ax2.set_ylim(bottom=-slide_detector_max, top=slide_detector_max)
ax2.set_ylabel('-10 box PWM score', color='purple')
ax2.tick_params(axis='y', labelcolor='purple')


# ---------- Adjust plot cosmetics ----------
ax1.ticklabel_format(style='plain', axis='x', useOffset=False)
ax1.set_xticks(np.arange(x_start, x_end+1, 500))
ax1.set_xticklabels(np.arange(x_start, x_end+1, 500), rotation=45)
ax1.set_xlabel('Position')
plt.title('Strand-specific coverage across genetic elements')

# Combine legends from both axes
handles1, labels1 = ax1.get_legend_handles_labels()
handles2, labels2 = ax2.get_legend_handles_labels()
all_handles = handles1 + handles2
all_labels = labels1 + labels2

# Deduplicate legend entries
unique_labels = dict(zip(all_labels, all_handles))
ax1.legend(unique_labels.values(), unique_labels.keys(), loc='upper right').set_zorder(1000)



plt.tight_layout()
plt.show()

Make legend appear at the side of the graph, instead of on top.