import numpy as np

class RamachandranSampler:
    def __init__(self):
        # the weights tend towards the beta strand region to avoid collapse

        # generalised for 18 amino acids
        self.general_dist = [
            # beta strand (top left)
            {'phi_m': -120, 'phi_s': 20, 'psi_m': 140, 'psi_s': 20, 'w': 0.70},
            # right handed helix (bottom left)
            {'phi_m': -65,  'phi_s': 15, 'psi_m': -40, 'psi_s': 15, 'w': 0.25},
            # left handed helix (top right)
            {'phi_m': 60,   'phi_s': 15, 'psi_m': 40,  'psi_s': 15, 'w': 0.05}
        ]

        # glycine more flexible
        self.glycine_dist = [
            # beta strand (top left)
            {'phi_m': -100, 'phi_s': 30, 'psi_m': 140, 'psi_s': 30, 'w': 0.3},
            # right handed helix (bottom left)
            {'phi_m': -60,  'phi_s': 30, 'psi_m': -30, 'psi_s': 30, 'w': 0.2},
            # left handed helix (top right)
            {'phi_m': 60,   'phi_s': 30, 'psi_m': 30,  'psi_s': 30, 'w': 0.2},
            # bottom right for glycine
            {'phi_m': 100,  'phi_s': 30, 'psi_m': -140,'psi_s': 30, 'w': 0.3}
        ]

        # proline no left handed helix
        self.proline_dist = [
            # beta strand (top left)
            {'phi_m': -63,  'phi_s': 5,  'psi_m': 150, 'psi_s': 20, 'w': 0.8},
            # right handed helix (bottom left)
            {'phi_m': -63,  'phi_s': 5,  'psi_m': -35, 'psi_s': 20, 'w': 0.2}
        ]

    # returns the correct distribution
    def _get_distribution(self, res_name):
        if res_name == 'GLY':
            return self.glycine_dist
        elif res_name == 'PRO':
            return self.proline_dist
        else:
            return self.general_dist

    # returns (phi, psi) for the residue (units degrees)
    def sample(self, res_name): 
        # the possible regions to sample 
        dist_options = self._get_distribution(res_name)
        
        weights = [d['w'] for d in dist_options]
        # normalise the weights to ensure sum is exactly 1.0 (no float errors)
        weights = np.array(weights) / np.sum(weights)
        # choose a region based on the weights 
        choice_idx = np.random.choice(len(dist_options), p=weights)
        selected = dist_options[choice_idx]
        
        # do a gaussian sample to get some variation
        phi = np.random.normal(selected['phi_m'], selected['phi_s'])
        psi = np.random.normal(selected['psi_m'], selected['psi_s'])
        
        # wrap angles to +-180 degrees
        phi = (phi + 180) % 360 - 180
        psi = (psi + 180) % 360 - 180
        
        return phi, psi

# example
if __name__ == "__main__":
    sampler = RamachandranSampler()
    
    print("Sampling 5 residues...")
    for res in ['ALA', 'GLY', 'PRO', 'TRP', 'VAL']:
        phi, psi = sampler.sample(res)
        print(f"{res}: Phi={phi:.1f}, Psi={psi:.1f}")
